Unable to return `int32` output
- Dominant language
- Python
- Stars
- 5.4k
- Forks
- 850
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 10
Description
## πDescribing the bug
It seems that has non `float32` outputs will always get cast to `float32`. This breaks segmentation models that return per-pixel indices. Here's the relevant source code:
https://github.com/apple/coremltools/blob/973eae67f2f273a29e80a9b009987516a070a58b/coremltools/converters/mil/backend/nn/passes/alert_return_type_cast.py#L13-L20
This contrasts with the fact that CoreML [does support `Int32` return types](https://developer.apple.com/documentation/coreml/mlmultiarray?language=objc). And this is further contrasted by the fact that `coremltools` itself supports `int32` outputs:
https://github.com/apple/coremltools/blob/973eae67f2f273a29e80a9b009987516a070a58b/coremltools/converters/mil/backend/mil/passes/adjust_io_to_supported_types.py#L126-L133
Why the discrepancy?
## To Reproduce
Run:
```python
from coremltools import convert, TensorType
from torch import int32, randn
from torch.jit import trace
from torch.nn import Module
class Model (Module):
def __init__ (self):
super().__init__()
def forward (self, input0):
return input0.to(int32)
example_input = randn(1, 256, 256, 3)
model = Model()
scripted_model = trace(model, [example_input])
coreml_model = convert(scripted_model, inputs=[TensorType(shape=example_input.shape)])
```
And observe:
```
Running MIL Common passes: 100%|βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ| 37/37 [00:00<00:00, 616.52 passes/s]
Running MIL Clean up passes: 100%|ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ| 9/9 [00:00<00:00, 8418.54 passes/s]
WARNING:root:Output var var_6 of type int32 in function main is cast to type fp32
```
## System environment (please complete the following information):
- coremltools version: 5.2.0 and 6.0b1
- OS (e.g. MacOS version or Linux type): macOS 12.3.1
- Any other relevant version information (e.g. PyTorch or TensorFlow version): Torch 1.11
Contributor guide
Research direction
Start with the linked alert_return_type_cast.py pass and compare it with adjust_io_to_supported_types.py, especially the referenced lines. Run the provided PyTorch tracing and conversion example to reproduce the warning, then inspect existing converter tests for output types. Done means supported int32 outputs remain int32 without the cast warning.
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
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100