Torchscript to CoreML conversion skips for loops
- Dominant language
- Python
- Stars
- 5.4k
- Forks
- 850
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 10
Description
## 🐞Description
When torchscript code contains for loops and is converted to coreML, the for loops are skipped and the result before the for loop is returned.
## To Reproduce
```
import numpy as np
import torch
import torch.nn as nn
import coremltools as ct
@torch.jit.script
def experiment(val: torch.Tensor):
val = int(val.item())
result = torch.zeros(val)
for i in range(val):
result[i] = i
return result
class Experiment(nn.Module):
def __init__(self):
super(Experiment, self).__init__()
def forward(self, x):
return experiment(x)
exp = Experiment()
# Use a tensor as input, not an integer
input_tensor = torch.tensor(100)
output = exp(input_tensor)
print(output)
traced_exp = torch.jit.script(exp, input_tensor)
traced_exp.eval()
output_traced = traced_exp(input_tensor)
print(output)
# Specify the input type as ct.TensorType(name="x", shape=(1,))
coreml_exp = ct.convert(
traced_exp,
source="pytorch",
inputs=[ct.TensorType(name="x", shape=(1,))],
convert_to="mlprogram"
)
# Create an input dictionary with the necessary input data
input_data = {
'x': np.array([100.0]),
}
# Make a prediction using the model
coreml_output = coreml_exp.predict(input_data)
print(coreml_output)
```
- Output
```
tensor([ 0., 1., 2., 3., 4., 5., 6., 7., 8., 9., 10., 11., 12., 13.,
14., 15., 16., 17., 18., 19., 20., 21., 22., 23., 24., 25., 26., 27.,
28., 29., 30., 31., 32., 33., 34., 35., 36., 37., 38., 39., 40., 41.,
42., 43., 44., 45., 46., 47., 48., 49., 50., 51., 52., 53., 54., 55.,
56., 57., 58., 59., 60., 61., 62., 63., 64., 65., 66., 67., 68., 69.,
70., 71., 72., 73., 74., 75., 76., 77., 78., 79., 80., 81., 82., 83.,
84., 85., 86., 87., 88., 89., 90., 91., 92., 93., 94., 95., 96., 97.,
98., 99.])
tensor([ 0., 1., 2., 3., 4., 5., 6., 7., 8., 9., 10., 11., 12., 13.,
14., 15., 16., 17., 18., 19., 20., 21., 22., 23., 24., 25., 26., 27.,
28., 29., 30., 31., 32., 33., 34., 35., 36., 37., 38., 39., 40., 41.,
42., 43., 44., 45., 46., 47., 48., 49., 50., 51., 52., 53., 54., 55.,
56., 57., 58., 59., 60., 61., 62., 63., 64., 65., 66., 67., 68., 69.,
70., 71., 72., 73., 74., 75., 76., 77., 78., 79., 80., 81., 82., 83.,
84., 85., 86., 87., 88., 89., 90., 91., 92., 93., 94., 95., 96., 97.,
98., 99.])
Converting PyTorch Frontend ==> MIL Ops: 0%| | 0/9 [00:00 MIL Ops: 78%|███████▊ | 7/9 [00:00<00:00, 3895.47 ops/s]
Running MIL frontend_pytorch pipeline: 100%|██████████| 5/5 [00:00<00:00, 23275.83 passes/s]
Running MIL default pipeline: 0%| | 0/66 [00:00
Contributor guide
Assessment
This issue has not been assessed yet.