apple / apple/coremltools

Exporting ct.convert crashes for torch.unfold (im2col) with flexible shapes

Open
#2,446 2 comments 0 reactions 0 assignees View on GitHub
bug Flexible Shape triaged
Dominant language
Python
Stars
5.4k
Forks
850
Avg merge
4d 5h
Merged PRs (30d)
10

Description

## Description
The torch.Unfold op conversion (im2col) crashes with flexible shaped input tensors.

## Stack Trace
```
ERROR - converting 'im2col' op (located at: '14'):

Converting PyTorch Frontend ==> MIL Ops: 92%|█████████▏| 12/13 [00:00<00:00, 563.92 ops/s]
Traceback (most recent call last):
File "playground.py", line 375, in
coreml_model = ct.convert(
File "/lib/python3.10/site-packages/coremltools/converters/_converters_entry.py", line 635, in convert
mlmodel = mil_convert(
File "/lib/python3.10/site-packages/coremltools/converters/mil/converter.py", line 186, in mil_convert
return _mil_convert(
File "/lib/python3.10/site-packages/coremltools/converters/mil/converter.py", line 218, in _mil_convert
proto, mil_program = mil_convert_to_proto(
File "/lib/python3.10/site-packages/coremltools/converters/mil/converter.py", line 294, in mil_convert_to_proto
prog = frontend_converter(model, **kwargs)
File "/lib/python3.10/site-packages/coremltools/converters/mil/converter.py", line 106, in __call__
return load(*args, **kwargs)
File "/lib/python3.10/site-packages/coremltools/converters/mil/frontend/torch/load.py", line 88, in load
return _perform_torch_convert(converter, debug)
File "/lib/python3.10/site-packages/coremltools/converters/mil/frontend/torch/load.py", line 151, in _perform_torch_convert
prog = converter.convert()
File "/lib/python3.10/site-packages/coremltools/converters/mil/frontend/torch/converter.py", line 1380, in convert
convert_nodes(self.context, self.graph, early_exit=not has_states)
File "/lib/python3.10/site-packages/coremltools/converters/mil/frontend/torch/ops.py", line 117, in convert_nodes
raise e # re-raise exception
File "/lib/python3.10/site-packages/coremltools/converters/mil/frontend/torch/ops.py", line 112, in convert_nodes
convert_single_node(context, node)
File "/lib/python3.10/site-packages/coremltools/converters/mil/frontend/torch/ops.py", line 173, in convert_single_node
add_op(context, node)
File "/lib/python3.10/site-packages/coremltools/converters/mil/frontend/torch/ops.py", line 8142, in im2col
indices = _construct_unfold_indices(N, C, H, W, kernel_size, stride)
File "/lib/python3.10/site-packages/coremltools/converters/mil/frontend/torch/ops.py", line 8060, in _construct_unfold_indices
offset_idx = np.arange(0, row_extent, stride[0])[None, :, None] * W + np.arange(0, col_extent, stride[1])
File "/lib/python3.10/site-packages/sympy/core/expr.py", line 340, in __float__
raise TypeError("Cannot convert expression to float")
TypeError: Cannot convert expression to float
```
## To Reproduce
Working example with fixed shape input.
```
import torch
import coremltools as ct
from torch import nn
import torch.nn.functional as F

class UnfoldModule(nn.Module):
def forward(self, x):
return F.unfold(x, (3, 3), padding=0)

model = UnfoldModule()
x = torch.rand(1, 1, 64, 64)

inpt = (x,)
exported_model = torch.export.export(model, inpt)
traced_model = torch.jit.trace(model, inpt)

ct.convert(traced_model, inputs=[ct.TensorType(shape=x.shape),],).save("/tmp/model.mlpackage")
```

Non-working example with flexible shape input.
```
import torch
import coremltools as ct
from torch import nn
import torch.nn.functional as F

class UnfoldModule(nn.Module):
def forward(self, x):
return F.unfold(x, (3, 3), padding=0)

model = UnfoldModule()
x = torch.rand(1, 1, 64, 64)

inpt = (x,)
exported_model = torch.export.export(model, inpt)
traced_model = torch.jit.trace(model, inpt)

flexible_shape = ct.Shape(shape=(1, 1, ct.RangeDim(lower_bound=64, upper_bound=128, default=64),
ct.RangeDim(lower_bound=64, upper_bound=128, default=64)))

ct.convert(traced_model, inputs=[ct.TensorType(shape=flexible_shape),],).save("/tmp/model.mlpackage")
```

## System environment
Ubuntu 22.04.5 LTS
Python 3.10.12
torch 2.3.1
torchvision 0.18.1
coremltools 8.2

Contributor guide

Open the contributing guide

Research direction

Start in coremltools/converters/mil/frontend/torch/ops.py at im2col and _construct_unfold_indices, then reproduce the flexible-shape example from the issue. Compare it with the fixed-shape conversion and add or run a regression test covering RangeDim inputs; done means flexible torch.unfold conversion completes without the SymPy float-conversion crash.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
68/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.