Cannot convert a model with torch.jit.script

[Ykid]

❓Question

I'm trying to convert a simple model like the one provided as follows:

import coremltools
import torch
from torch import nn
import numpy as np

class ControlFlowNet3(nn.Module):

    def forward(self, flag_tensor: torch.Tensor):
        if flag_tensor[0] == 1:
            return torch.tensor([0]), torch.tensor([1])
        else:
            return torch.tensor([1]), torch.tensor([2])


def main():
    model = ControlFlowNet3()
    model.eval()
    scripted_model = torch.jit.script(model)
    flag_tensor = coremltools.TensorType(name='flag_tensor', shape=(1,), dtype=np.int32)
    mlmodel = coremltools.converters.convert(
        scripted_model,
        inputs=[
            flag_tensor
        ],
        minimum_deployment_target=coremltools.target.iOS15
    )
    print('done')


if __name__ == '__main__':
    main()

but I ran into an error

Support for converting Torch Script Models is experimental. If possible you should use a traced model for conversion.
Converting PyTorch Frontend ==> MIL Ops:   0%|          | 0/9 [00:00<?, ? ops/s]
Converting PyTorch Frontend ==> MIL Ops:  80%|████████  | 4/5 [00:00<00:00, 14388.69 ops/s]
Converting PyTorch Frontend ==> MIL Ops:  89%|████████▉ | 8/9 [00:00<00:00, 3460.29 ops/s]
Traceback (most recent call last):
  File "/Users/abc/Desktop/code/integration/Packages/cde/notebooks/scripts/debug_script.py", line 31, in <module>
    main()
  File "/Users/abc/Desktop/code/integration/Packages/cde/notebooks/scripts/debug_script.py", line 20, in main
    mlmodel = coremltools.converters.convert(
  File "/Users/abc/Library/Caches/pypoetry/virtualenvs/notebooks-tDop3foe-py3.10/lib/python3.10/site-packages/coremltools/converters/_converters_entry.py", line 551, in convert
    mlmodel = mil_convert(
  File "/Users/abc/Library/Caches/pypoetry/virtualenvs/notebooks-tDop3foe-py3.10/lib/python3.10/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/abc/Library/Caches/pypoetry/virtualenvs/notebooks-tDop3foe-py3.10/lib/python3.10/site-packages/coremltools/converters/mil/converter.py", line 212, in _mil_convert
    proto, mil_program = mil_convert_to_proto(
  File "/Users/abc/Library/Caches/pypoetry/virtualenvs/notebooks-tDop3foe-py3.10/lib/python3.10/site-packages/coremltools/converters/mil/converter.py", line 286, in mil_convert_to_proto
    prog = frontend_converter(model, **kwargs)
  File "/Users/abc/Library/Caches/pypoetry/virtualenvs/notebooks-tDop3foe-py3.10/lib/python3.10/site-packages/coremltools/converters/mil/converter.py", line 108, in __call__
    return load(*args, **kwargs)
  File "/Users/abc/Library/Caches/pypoetry/virtualenvs/notebooks-tDop3foe-py3.10/lib/python3.10/site-packages/coremltools/converters/mil/frontend/torch/load.py", line 75, in load
    return _perform_torch_convert(converter, debug)
  File "/Users/abc/Library/Caches/pypoetry/virtualenvs/notebooks-tDop3foe-py3.10/lib/python3.10/site-packages/coremltools/converters/mil/frontend/torch/load.py", line 114, in _perform_torch_convert
    prog = converter.convert()
  File "/Users/abc/Library/Caches/pypoetry/virtualenvs/notebooks-tDop3foe-py3.10/lib/python3.10/site-packages/coremltools/converters/mil/frontend/torch/converter.py", line 484, in convert
    convert_nodes(self.context, self.graph)
  File "/Users/abc/Library/Caches/pypoetry/virtualenvs/notebooks-tDop3foe-py3.10/lib/python3.10/site-packages/coremltools/converters/mil/frontend/torch/ops.py", line 93, in convert_nodes
    add_op(context, node)
  File "/Users/abc/Library/Caches/pypoetry/virtualenvs/notebooks-tDop3foe-py3.10/lib/python3.10/site-packages/coremltools/converters/mil/frontend/torch/ops.py", line 3271, in _if
    cond = mb.cond(
  File "/Users/abc/Library/Caches/pypoetry/virtualenvs/notebooks-tDop3foe-py3.10/lib/python3.10/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/abc/Library/Caches/pypoetry/virtualenvs/notebooks-tDop3foe-py3.10/lib/python3.10/site-packages/coremltools/converters/mil/mil/builder.py", line 181, in _add_op
    new_op.build_nested_blocks()
  File "/Users/abc/Library/Caches/pypoetry/virtualenvs/notebooks-tDop3foe-py3.10/lib/python3.10/site-packages/coremltools/converters/mil/mil/ops/defs/iOS15/control_flow.py", line 103, in build_nested_blocks
    true_block.set_outputs(true_ret_vars)
  File "/Users/abc/Library/Caches/pypoetry/virtualenvs/notebooks-tDop3foe-py3.10/lib/python3.10/site-packages/coremltools/converters/mil/mil/block.py", line 294, in set_outputs
    raise ValueError(msg.format(ov.name, self.name, self))
AttributeError: 'tuple' object has no attribute 'name'

wondering what could be done to fix the issue?

Thank you!

[Ykid]

versions

torch==2.0.0
coremltools==7.0
numpy==1.24.3

[TobyRoseman]

This is related #1995.

Note the first line of your output:

Support for converting Torch Script Models is experimental. If possible you should use a traced model for conversion.

If you call torch.jit.trace on your model prior to conversion, it should work fine.

[Ykid]

Thanks @TobyRoseman for the reply! If I understand it correctly, the traced model will not contain any control flow logic. So in my case torch will only trace one of the if / else branches right?

[TobyRoseman]

@Ykid - I believe that is correct. However I am not a PyTorch expert. I suggest asking that question in a PyTorch forum. If you do that, please share a link to it here.