Can not convert Pytorch model to CoreML

[pkarimib]

Error Description

Error when running the example in Model-Scripting. I can not convert the Pytorch model to CoreML.

Trace

Traceback (most recent call last):
  File "/Users/panteababaahmadi/Documents/GitHub/Joanna_code/first_order_model/coreml.py", line 311, in <module>
    mlmodel = coremltools.converters.convert(
  File "/usr/local/lib/python3.9/site-packages/coremltools/converters/_converters_entry.py", line 326, in convert
    mlmodel = mil_convert(
  File "/usr/local/lib/python3.9/site-packages/coremltools/converters/mil/converter.py", line 182, in mil_convert
    return _mil_convert(model, convert_from, convert_to, ConverterRegistry, MLModel, compute_units, **kwargs)
  File "/usr/local/lib/python3.9/site-packages/coremltools/converters/mil/converter.py", line 209, in _mil_convert
    proto, mil_program = mil_convert_to_proto(
  File "/usr/local/lib/python3.9/site-packages/coremltools/converters/mil/converter.py", line 300, in mil_convert_to_proto
    prog = frontend_converter(model, **kwargs)
  File "/usr/local/lib/python3.9/site-packages/coremltools/converters/mil/converter.py", line 104, in __call__
    return load(*args, **kwargs)
  File "/usr/local/lib/python3.9/site-packages/coremltools/converters/mil/frontend/torch/load.py", line 49, in load
    converter = TorchConverter(torchscript, inputs, outputs, cut_at_symbols)
  File "/usr/local/lib/python3.9/site-packages/coremltools/converters/mil/frontend/torch/converter.py", line 150, in __init__
    raw_graph, params_dict = self._expand_and_optimize_ir(self.torchscript)
  File "/usr/local/lib/python3.9/site-packages/coremltools/converters/mil/frontend/torch/converter.py", line 454, in _expand_and_optimize_ir
    graph, params_dict = TorchConverter._jit_pass_lower_graph(graph, torchscript)
  File "/usr/local/lib/python3.9/site-packages/coremltools/converters/mil/frontend/torch/converter.py", line 390, in _jit_pass_lower_graph
    _lower_graph_block(graph)
  File "/usr/local/lib/python3.9/site-packages/coremltools/converters/mil/frontend/torch/converter.py", line 355, in _lower_graph_block
    _lower_graph_block(block)
  File "/usr/local/lib/python3.9/site-packages/coremltools/converters/mil/frontend/torch/converter.py", line 376, in _lower_graph_block
    is_tensor = _check_is_tensor(node, module)
  File "/usr/local/lib/python3.9/site-packages/coremltools/converters/mil/frontend/torch/converter.py", line 333, in _check_is_tensor
    assert str(node.output().type()) == "Tensor"
AssertionError

To Reproduce

Here is the python script to produce the error:

import torch
import torch.nn as nn
import torch.nn.functional as F
import coremltools

class _LoopBody(nn.Module):
    def __init__(self, channels):
        super(_LoopBody, self).__init__()
        conv = nn.Conv2d(
            in_channels=channels,
            out_channels=channels,
            kernel_size=3,
            padding=1,
        )
        self.conv = conv

    def forward(self, x):
        x = self.conv(x)
        x = F.relu(x)
        return x

class ControlFlowNet(nn.Module):
    def __init__(self, num_channels: int):
        super(ControlFlowNet, self).__init__()
        self.loop_body = _LoopBody(num_channels)
        print("here 1")

    def forward(self, x):
        avg = torch.mean(x)
        if avg.item() < 0:
            loop_count = 2
        else:
            loop_count = 1
        print("here 2")
        for _ in range(loop_count):
            x = self.loop_body(x)
        return x


model = ControlFlowNet(num_channels=3)
scripted_model = torch.jit.script(model)

mlmodel = coremltools.converters.convert(
  scripted_model,
  inputs=[coremltools.TensorType(shape=(1, 3, 64, 64))],
)
mlmodel.save("a.mlmodel")

System environment

• coremltools version: 5.1.0
• OS: MacOS
• macOS version: 10.15.7
• How you install python: system
• python version: 3.9
• torch version: 1.9

[trulyspinach]

Can confirm the same issue using 5.1 and torch 1.9

[JRGit4UE]

The issue seems to be a bias set as Optional which asserts in CoreMLTools converter.py. node.output.type== Optional[Tensor] 19 defined in (%19 : Tensor? = prim::GetAttr[name="bias"](%17) Here is the TorchScript graph:

here 1
graph(%self : __torch__.ControlFlowNet,
      %x.1 : Tensor):
  %20 : bool = prim::Constant[value=1]() # C:/yolact/test-coremltools.py:35:8
  %16 : str = prim::Constant[value="here 2"]() # C:/yolact/test-coremltools.py:34:14
  %3 : NoneType = prim::Constant()
  %7 : int = prim::Constant[value=0]() # C:/yolact/test-coremltools.py:30:24
  %10 : int = prim::Constant[value=2]() # C:/yolact/test-coremltools.py:31:25
  %11 : int = prim::Constant[value=1]() # C:/yolact/test-coremltools.py:33:25
  %avg.1 : Tensor = aten::mean(%x.1, %3) # C:/yolact/test-coremltools.py:29:14
  %6 : Scalar = aten::item(%avg.1) # C:/yolact/test-coremltools.py:30:11
  %8 : bool = aten::lt(%6, %7) # C:/yolact/test-coremltools.py:30:11
  %loop_count : int = prim::If(%8) # C:/yolact/test-coremltools.py:30:8
    block0():
      -> (%10)
    block1():
      -> (%11)
   = prim::Print(%16) # C:/yolact/test-coremltools.py:34:8
  %x : Tensor = prim::Loop(%loop_count, %20, %x.1) # C:/yolact/test-coremltools.py:35:8
    block0(%21 : int, %x.13 : Tensor):
      %22 : __torch__._LoopBody = prim::GetAttr[name="loop_body"](%self)
      %x.7 : Tensor = prim::CallMethod[name="forward"](%22, %x.13) # C:/yolact/test-coremltools.py:36:16
      -> (%20, %x.7)
  return (%x)

Are there any ideas for a quick workaround?
Should I change the type of the bias from Optional to an empty (zero_) Tensor in the TorchScript model?
How can I find the location %17 or %19 mentioned above in the assertion message, it's not appearing in the graph?

[JRGit4UE]

The workaround seems to be to define loop_count as

loop_count:int=0 before it is used ..

[GaganNarula]

Hi

I faced the same problem on my mac with the latest beta release of coremltools (6.0b2) while trying to convert a model that uses torchaudio.pipelines.HUBERT_LARGE as a pre-trained backbone.

Environment torch 1.11.0 torchaudio 0.11.0 coremltools 6.0b2 OSX Monterrey (Mac M1 architecture) python 3.9

All packages installed with pip in a conda environment.

My work-around for this problem was to replace the assert for "Tensor" in the _jit_pass_lower_graph._check_is_tensor() method of TorchConverter class in coremltools/converters/mil/frontend/torch/converter.py with the following lines

type_ = str(node.output().type())
assert type_ == "Tensor" or type_ == "Optional[Tensor]"

Its a hack but it works for now. Model compiles without errors.

[TobyRoseman]

Our support for Torch Script Models is experimental. If at all possible trace your model prior to conversion.

[JRGit4UE]

@TobyRoseman Thank you for your effort improving CoreMLTools, I really appreciate it, but I would suggest that your latest comment about experimental support of Torch Script Models should pop up somewhere in the release notes, e.g. for version 6.0 as well, so that we know what to expect in that case...

[TobyRoseman]

@JRGit4UE - Thanks for the feedback. I agree it would be good to make this more clear. Users do get a warning when they try to convert a Torch Script Model. So that's good. This isn't something new to the 6.0 release. So I don't think it makes sense to mention it in the release notes. Let me know if you have other ideas.

[TobyRoseman]

This is a duplicate of #1504. We have an open pull request (#1657).

Closing this issue as a duplicate.