apple / apple/coremltools

Coremltools inconsistency between torch.jit.script and torch.jit.trace conversion

Open
#1,400 3 comments 1 reaction 0 assignees View on GitHub
bug PyTorch (not traced) PyTorch (traced) triaged
Dominant language
Python
Stars
5.4k
Forks
850
Avg merge
4d 5h
Merged PRs (30d)
10

Description

## 🐞Describe the bug

Converting a simple test model in Pytorch with only a Conv2d() operation in its forward() method,
the conversion succeeds with **torch.jit.trace** but fails with **torch.jit.script**, which seems to be not consistent.

## Trace
If applicable, please paste the error trace.
Trying to convert the following example fails with an assertion
```
assert str(node.output().type()) == "Tensor"
```
in coremltools\converters\mil\frontend\torch\converter.py function '_check_is_tensor()'
for parameter **bias** in Conv2d() which is defined as **Optional** (`c10::optional& bias_opt`) in
https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/Convolution.cpp
```
at::Tensor conv2d(
const Tensor& input_, const Tensor& weight, const c10::optional& bias_opt,
IntArrayRef stride, IntArrayRef padding, IntArrayRef dilation, int64_t groups)
```
When compiling the model as **script**
function _check_is_tensor() runs into an assertion, as **Optional[Tensor] != Tensor**
This error does not occur when compiling the model via **trace**, but that is not what I need.

## To Reproduce
- If a python script can reproduce the error, please paste the code snippet
```
class Test(nn.Module):
def __init__(self):
super().__init__()
self.conv = nn.Conv2d(1,1,3, bias=True)
def forward(self, x:Tensor)->Tensor:
y = self.conv(x)
return y
```

```
import coremltools as ct
with torch.no_grad():
torch.set_default_tensor_type('torch.FloatTensor')
b = torch.rand(1,1,4,4)
t = Test()
t.eval()
sd = t.state_dict()
sd['conv.bias'] = torch.zeros(1)
t.load_state_dict(sd)
res = t(b)

script = torch.jit.script(t) # does not work
mlmodel = ct.convert(
script,
source="pytorch",
inputs=[ct.TensorType(name='x', shape=b.shape)],
convert_to="mlprogram",
debug=False
)
```

```
script = torch.jit.trace(t,b) # does work, but I would need torch.jit.script
mlmodel = ct.convert(
script,
source="pytorch",
inputs=[ct.TensorType(name='x', shape=b.shape)],
convert_to="mlpackage",
debug=False
)
```

Contributor guide

Open the contributing guide

Research direction

Start with coremltools/converters/mil/frontend/torch/converter.py, especially _check_is_tensor(), and reproduce the failure using the provided scripted Conv2d model. Compare the scripted and traced conversion paths and verify that the scripted model converts without the Optional[Tensor] assertion.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.