Issues converting a PyTorch model to CoreML

[eightysix]

Hello.
I’m an iOS/macOS developer who’s very new to AI.
I’d like to integrate to test some radiology-related AI models in a Swift environment, but I’m having issues converting them to CoreML.
Specifically, I’m encountering this error: AssertionError: tensor value not consistent between torch ir and state_dict, for which I’ve found no information anywhere.

This is my full Jupyter notebook for conversion:

import torchxrayvision as xrv
import skimage, torch, torchvision

model = xrv.models.DenseNet(weights="densenet121-res224-rsna")
model.eval()

example_input = torch.rand(1, 1, 224, 224)
traced_model = torch.jit.trace(model, example_input)
out = traced_model(example_input)

import coremltools as ct
scale = 1/(0.226*255.0)
bias = [- 0.485/(0.229) , - 0.456/(0.224), - 0.406/(0.225)]

image_input = ct.ImageType(name="input_1",
                           shape=example_input.shape,
                           scale=scale, bias=bias)

model = ct.convert(
    traced_model,
    inputs=[image_input],
    compute_units=ct.ComputeUnit.CPU_ONLY,
)

[Krass1]

Hi, I was able to successfully convert it, but had to make a few changes. The changes and the edited code are below. Hope it helps.

1. Used coremltools v7.0, torch v2.1.0, torchvision v0.16, python v3.10
2. coremltools needs a custom operation isnan to be added to the ops.py file. The function I came up with looks like this:

@register_torch_op
def isnan(context, node):
    torchTensor, = _get_inputs(context, node, expected=1)
    x = mb.const(val=(torchTensor != torchTensor), name=node.name)
    context.add(x, node.name)

3. The scale and bias need to be changed to match what the model you are using expects. In the case of densenet121-res224-rsna the scale is 1/1024 and bias is 0.
4. Lastly we have to add color_layout=ct.colorlayout.GRAYSCALE to the input_image in order to tell coreML to expect only one channel of color and not three.

Here is the code that works for me and generates a coreML model, but I have not put this model in an app to see if it performs as expected.

import torchxrayvision as xrv
import skimage, torch, torchvision
import coremltools as ct

model = xrv.models.DenseNet(weights="densenet121-res224-rsna")
model.eval()

example_input = torch.rand(1, 1, 1024, 1024)
traced_model = torch.jit.trace(model, example_input)

scale = 1/(1024)
bias = [0, 0, 0]

image_input = ct.ImageType(name="input_1", shape=example_input.shape, color_layout=ct.colorlayout.GRAYSCALE, scale=scale, bias=bias)
output = ct.TensorType(name="output")

mlmodel = ct.convert(traced_model, inputs=[image_input], outputs=[output])

print(mlmodel)```