apple / apple/coremltools

Unable to return `int32` output

Open
#1,546 4 comments 0 reactions 0 assignees View on GitHub
bug triaged
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.