Coremltools inconsistency between torch.jit.script and torch.jit.trace conversion
- 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
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