apple / apple/coremltools

Torchscript to CoreML conversion skips for loops

Open
#1,995 5 comments 0 reactions 0 assignees View on GitHub
bug PyTorch (not traced)
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.