apple / apple/coremltools

Complex op such as `irfftn` doesn't support dynamic shapes

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

Description

## 🐞Describing the bug
- Trying to convert this Pytorch [lama model ](https://github.com/mallman/CoreMLaMa) to CoreML, and it does work fine when the model uses Fixed Input size. However, when I want to use flexible shapes, you get shape mismatch error between real data and imag_data. Seems like it's not able to correctly process the symbolic shapes properly.

## Stack Trace
```
DEBUG:coremltools:Adding const op '256_end_0'
INFO:coremltools:Adding op '256_end_0' of type const
DEBUG:coremltools:Downcast const op 256_end_0 dataint64 as int32
DEBUG:coremltools:Downcast const op 256_end_0 dataint64 as int32
DEBUG:coremltools:Adding const op '256_end_mask_0'
INFO:coremltools:Adding op '256_end_mask_0' of type const
DEBUG:coremltools:Adding const op '256_squeeze_mask_0'
INFO:coremltools:Adding op '256_squeeze_mask_0' of type const
INFO:coremltools:Converting op ffted3.3 : complex
INFO:coremltools:Adding op 'complex_0' of type complex
Converting PyTorch Frontend ==> MIL Ops: 5%|███████▍ | 152/3297 [00:00<00:04, 650.29 ops/s]
Traceback (most recent call last):
File "/Users/philip/SkyCoreML/main.py", line 51, in
main()
File "/Users/philip/SkyCoreML/main.py", line 43, in main
lama.convert(args)
File "/Users/philip/SkyCoreML/src/skycoreml/conversion/lama.py", line 124, in convert
coreml_model = ct.convert(
File "/Users/philip/SkyCoreML/env/lib/python3.9/site-packages/coremltools/converters/_converters_entry.py", line 530, in convert
mlmodel = mil_convert(
File "/Users/philip/SkyCoreML/env/lib/python3.9/site-packages/coremltools/converters/mil/converter.py", line 188, in mil_convert
return _mil_convert(model, convert_from, convert_to, ConverterRegistry, MLModel, compute_units, **kwargs)
File "/Users/philip/SkyCoreML/env/lib/python3.9/site-packages/coremltools/converters/mil/converter.py", line 212, in _mil_convert
proto, mil_program = mil_convert_to_proto(
File "/Users/philip/SkyCoreML/env/lib/python3.9/site-packages/coremltools/converters/mil/converter.py", line 286, in mil_convert_to_proto
prog = frontend_converter(model, **kwargs)
File "/Users/philip/SkyCoreML/env/lib/python3.9/site-packages/coremltools/converters/mil/converter.py", line 108, in __call__
return load(*args, **kwargs)
File "/Users/philip/SkyCoreML/env/lib/python3.9/site-packages/coremltools/converters/mil/frontend/torch/load.py", line 63, in load
return _perform_torch_convert(converter, debug)
File "/Users/philip/SkyCoreML/env/lib/python3.9/site-packages/coremltools/converters/mil/frontend/torch/load.py", line 102, in _perform_torch_convert
prog = converter.convert()
File "/Users/philip/SkyCoreML/env/lib/python3.9/site-packages/coremltools/converters/mil/frontend/torch/converter.py", line 439, in convert
convert_nodes(self.context, self.graph)
File "/Users/philip/SkyCoreML/env/lib/python3.9/site-packages/coremltools/converters/mil/frontend/torch/ops.py", line 92, in convert_nodes
add_op(context, node)
File "/Users/philip/SkyCoreML/env/lib/python3.9/site-packages/coremltools/converters/mil/frontend/torch/ops.py", line 5691, in complex
result = mb.complex(real_data=real_part, imag_data=imag_part)
File "/Users/philip/SkyCoreML/env/lib/python3.9/site-packages/coremltools/converters/mil/mil/ops/registry.py", line 183, in add_op
return cls._add_op(op_cls_to_add, **kwargs)
File "/Users/philip/SkyCoreML/env/lib/python3.9/site-packages/coremltools/converters/mil/mil/builder.py", line 182, in _add_op
new_op.type_value_inference()
File "/Users/philip/SkyCoreML/env/lib/python3.9/site-packages/coremltools/converters/mil/mil/operation.py", line 253, in type_value_inference
output_types = self.type_inference()
File "/Users/philip/SkyCoreML/env/lib/python3.9/site-packages/coremltools/converters/mil/mil/ops/defs/complex_dialect_ops.py", line 162, in type_inference
raise ValueError(
ValueError: The shape of real_data ((1, 192, is32, is33)) and imag_data ((1, 192, is34, is35)) must match to construct complex data.
```

## To Reproduce

To reproduce, clone the above repository for CoreLama, and run `convert_lama.py`. The only thing you would want to change is giving some kind of flexible shape to the image input and mask input, so something like:

```
image_shape = ct.EnumeratedShapes(
shapes=[[1, 3, 256, 256], [1, 3, 512, 512], [1, 3, 1024, 1024]],
default=[1, 3, 512, 512],
)

mask_shape = ct.EnumeratedShapes(
shapes=[[1, 1, 256, 256], [1, 1, 512, 512], [1, 1, 1024, 1024]],
default=[1, 1, 512, 512],
)

coreml_model = ct.convert(
jit_model,
convert_to="mlprogram",
compute_precision=ct.precision.FLOAT32,
compute_units=ct.ComputeUnit.CPU_AND_GPU,
inputs=[
ct.ImageType(name="image",
shape=image_shape,
scale=1/255.0),
ct.ImageType(
name="mask",
shape=mask_shape,
color_layout=ct.colorlayout.GRAYSCALE)
],
outputs=[ct.ImageType(name="output")],
skip_model_load=True
)

```

## Things I have tried

Doesn't matter if you are tracing or scripting, it will fail as long as you enable flexible shape. It should produce mlpackage properly. I have tried multiple versions of Pytorch and CoreML tools, but no luck. Also tried going back to legacy neuralnetwork option, but below creates a model file that seems to have some kind of memory leak.

```
# see if you can update the spec
spec = ct.utils.load_spec(coreml_model_file_name)
image = spec.description.input[0].name
mask = spec.description.input[1].name

# update the image
flexible_shape_utils.set_multiarray_ndshape_range(
spec,
feature_name=image,
lower_bounds=[1, 3, 256, 256],
upper_bounds=[1, 3, 1024, 1024],
)

# update the mask
flexible_shape_utils.set_multiarray_ndshape_range(
spec,
feature_name=mask,
lower_bounds=[1, 1, 256, 256],
upper_bounds=[1, 1, 1024, 1024],
)

# save the spec
coreml_model_file_name = "LaMa_updated.mlmodel"
coreml_model_file_name = os.path.join(save_dir, coreml_model_file_name)
ct.utils.save_spec(spec, coreml_model_file_name)
```

## System environment (please complete the following information):
- coremltools version: 7.0b1
- OS (e.g. MacOS version or Linux type): Mac Mini with M2 chip, macOS Sonoma
- Any other relevant version information (e.g. PyTorch or TensorFlow version): Pytorch 2.0.1
- Tried both Python 3.9 and 3.11.

## Additional context
- I know that support for complex number is quite limited at the moment, but being able to use flexible shape is very important. Is there any way to work around this?

Contributor guide

Open the contributing guide

Research direction

Start by reproducing the failure with convert_lama.py using the flexible image and mask shapes. Read coremltools/converters/mil/frontend/torch/ops.py at the complex conversion and complex_dialect_ops.py at complex type inference. Done means the conversion accepts the symbolic real_data and imag_data shapes and produces the requested mlpackage.

Written by the indexing model from the issue text.

Assessment

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.