assert str(node.output().type()) == "Tensor" AssertionError during conversion of model from documentation site
🐞Describe the bug
I tried converting the ControlFlowNet example model from the documentation site (https://coremltools.readme.io/docs/model-scripting) using torch.jit.script() and coremltools.converters.convert(), but I saw the assert str(node.output().type()) == "Tensor" assertion error from the /site-packages/coremltools/converters/mil/frontend/torch/converter.py file during the process of converting the TorchScript object to a ML Model. I encountered this error on both Linux (Ubuntu 20.04.4 LTS) and Windows (Windows 10).
Stack Trace
- Stack Trace On Linux
Traceback (most recent call last):
File "coreml_example.py", line 40, in <module>
mlmodel = coremltools.converters.convert(
File "/home/jacky/.local/lib/python3.8/site-packages/coremltools/converters/_converters_entry.py", line 352, in convert
mlmodel = mil_convert(
File "/home/jacky/.local/lib/python3.8/site-packages/coremltools/converters/mil/converter.py", line 183, in mil_convert
return _mil_convert(model, convert_from, convert_to, ConverterRegistry, MLModel, compute_units, **kwargs)
File "/home/jacky/.local/lib/python3.8/site-packages/coremltools/converters/mil/converter.py", line 210, in _mil_convert
proto, mil_program = mil_convert_to_proto(
File "/home/jacky/.local/lib/python3.8/site-packages/coremltools/converters/mil/converter.py", line 273, in mil_convert_to_proto
prog = frontend_converter(model, **kwargs)
File "/home/jacky/.local/lib/python3.8/site-packages/coremltools/converters/mil/converter.py", line 105, in __call__
return load(*args, **kwargs)
File "/home/jacky/.local/lib/python3.8/site-packages/coremltools/converters/mil/frontend/torch/load.py", line 46, in load
converter = TorchConverter(torchscript, inputs, outputs, cut_at_symbols)
File "/home/jacky/.local/lib/python3.8/site-packages/coremltools/converters/mil/frontend/torch/converter.py", line 156, in __init__
raw_graph, params_dict = self._expand_and_optimize_ir(self.torchscript)
File "/home/jacky/.local/lib/python3.8/site-packages/coremltools/converters/mil/frontend/torch/converter.py", line 456, in _expand_and_optimize_ir
graph, params_dict = TorchConverter._jit_pass_lower_graph(graph, torchscript)
File "/home/jacky/.local/lib/python3.8/site-packages/coremltools/converters/mil/frontend/torch/converter.py", line 400, in _jit_pass_lower_graph
_lower_graph_block(graph)
File "/home/jacky/.local/lib/python3.8/site-packages/coremltools/converters/mil/frontend/torch/converter.py", line 365, in _lower_graph_block
_lower_graph_block(block)
File "/home/jacky/.local/lib/python3.8/site-packages/coremltools/converters/mil/frontend/torch/converter.py", line 386, in _lower_graph_block
is_tensor = _check_is_tensor(node, module)
File "/home/jacky/.local/lib/python3.8/site-packages/coremltools/converters/mil/frontend/torch/converter.py", line 343, in _check_is_tensor
assert str(node.output().type()) == "Tensor"
AssertionError
- Stack Trace On Windows
Traceback (most recent call last):
File "C:\Users\jacky\Desktop\coreml_test\coreml_example.py", line 40, in <module>
mlmodel = coremltools.converters.convert(
File "C:\Users\jacky\AppData\Local\Programs\Python\Python39\lib\site-packages\coremltools\converters\_converters_entry.py", line 352, in convert
mlmodel = mil_convert(
File "C:\Users\jacky\AppData\Local\Programs\Python\Python39\lib\site-packages\coremltools\converters\mil\converter.py", line 183, in mil_convert
return _mil_convert(model, convert_from, convert_to, ConverterRegistry, MLModel, compute_units, **kwargs)
File "C:\Users\jacky\AppData\Local\Programs\Python\Python39\lib\site-packages\coremltools\converters\mil\converter.py", line 210, in _mil_convert
proto, mil_program = mil_convert_to_proto(
File "C:\Users\jacky\AppData\Local\Programs\Python\Python39\lib\site-packages\coremltools\converters\mil\converter.py", line 273, in mil_convert_to_proto
prog = frontend_converter(model, **kwargs)
File "C:\Users\jacky\AppData\Local\Programs\Python\Python39\lib\site-packages\coremltools\converters\mil\converter.py", line 105, in __call__
return load(*args, **kwargs)
File "C:\Users\jacky\AppData\Local\Programs\Python\Python39\lib\site-packages\coremltools\converters\mil\frontend\torch\load.py", line 46, in load
converter = TorchConverter(torchscript, inputs, outputs, cut_at_symbols)
File "C:\Users\jacky\AppData\Local\Programs\Python\Python39\lib\site-packages\coremltools\converters\mil\frontend\torch\converter.py", line 156, in __init__
raw_graph, params_dict = self._expand_and_optimize_ir(self.torchscript)
File "C:\Users\jacky\AppData\Local\Programs\Python\Python39\lib\site-packages\coremltools\converters\mil\frontend\torch\converter.py", line 456, in _expand_and_optimize_ir
graph, params_dict = TorchConverter._jit_pass_lower_graph(graph, torchscript)
File "C:\Users\jacky\AppData\Local\Programs\Python\Python39\lib\site-packages\coremltools\converters\mil\frontend\torch\converter.py", line 400, in _jit_pass_lower_graph
_lower_graph_block(graph)
File "C:\Users\jacky\AppData\Local\Programs\Python\Python39\lib\site-packages\coremltools\converters\mil\frontend\torch\converter.py", line 365, in _lower_graph_block
_lower_graph_block(block)
File "C:\Users\jacky\AppData\Local\Programs\Python\Python39\lib\site-packages\coremltools\converters\mil\frontend\torch\converter.py", line 386, in _lower_graph_block
is_tensor = _check_is_tensor(node, module)
File "C:\Users\jacky\AppData\Local\Programs\Python\Python39\lib\site-packages\coremltools\converters\mil\frontend\torch\converter.py", line 343, in _check_is_tensor
assert str(node.output().type()) == "Tensor"
AssertionError
To Reproduce
- Please run the following code.
import torch
import torch.nn as nn
import torch.nn.functional as F
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)
def forward(self, x):
avg = torch.mean(x)
if avg.item() < 0:
loop_count = 2
else:
loop_count = 1
for _ in range(loop_count):
x = self.loop_body(x)
return x
model = ControlFlowNet(num_channels=3)
scripted_model = torch.jit.script(model)
import coremltools
mlmodel = coremltools.converters.convert(
scripted_model,
inputs=[coremltools.TensorType(shape=(1, 3, 64, 64))],
)
System environment (please complete the following information):
Linux
Distributor ID: Ubuntu
Description: Ubuntu 20.04.4 LTS
Release: 20.04
Codename: focal
------------------------------------------------
coremltools 5.2.0
torch 1.10.2+cu113
protobuf 3.19.3
packaging 21.3
sympy 1.10.1
numpy 1.22.4
tqdm 4.64.0
Windows
OS Name: Microsoft Windows 10 Home
Version: 10.0.19044 Build 19044
------------------------------------------------
coremltools 5.2.0
torch 1.10.2+cu113
protobuf 3.19.3
packaging 21.3
sympy 1.10.1
numpy 1.22.3
tqdm 4.62.3
Is there any update on this by any chance? Thanks a lot in advance!
Is there any update on this by any chance? Thanks a lot in advance!
Sorry, no update. Our support for Torch Script models is experimental. We have much better support for Torch Traced models.
Thank you for the detailed issue. I can reproduce the problem (in coremltools 6.0b1).