Basic torch.nn.LSTM can be converted from traced model but not scripted model

[simgt]

Trying to convert an unmodified torch.nn.LSTM model through tracing and scripting. Conversion of the traced model works, conversion of the scripted model fails. A self-contained script to reproduce the error is given bellow.

Trace

test_conversion.py:38: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
/home/vagrant/.local/lib/python3.6/site-packages/coremltools/converters/_converters_entry.py:182: in convert
    **kwargs
/home/vagrant/.local/lib/python3.6/site-packages/coremltools/converters/mil/converter.py:129: in mil_convert
    ConverterRegistry, **kwargs)
/home/vagrant/.local/lib/python3.6/site-packages/coremltools/converters/mil/converter.py:171: in mil_convert_to_proto
    prog = frontend_converter(model, **kwargs)
/home/vagrant/.local/lib/python3.6/site-packages/coremltools/converters/mil/converter.py:85: in __call__
    return load(*args, **kwargs)
/home/vagrant/.local/lib/python3.6/site-packages/coremltools/converters/mil/frontend/torch/load.py:70: in load
    converter = TorchConverter(torchscript, inputs, outputs, cut_at_symbols)
/home/vagrant/.local/lib/python3.6/site-packages/coremltools/converters/mil/frontend/torch/converter.py:145: in __init__
    raw_graph, params_dict = self._expand_and_optimize_ir(self.torchscript)
/home/vagrant/.local/lib/python3.6/site-packages/coremltools/converters/mil/frontend/torch/converter.py:263: in _expand_and_optimize_ir
    torchscript.forward.graph, torchscript._c
/home/vagrant/.local/lib/python3.6/site-packages/torch/jit/_script.py:558: in __getattr__
    return super(RecursiveScriptModule, self).__getattr__(attr)
/home/vagrant/.local/lib/python3.6/site-packages/torch/jit/_script.py:288: in __getattr__
    return super(ScriptModule, self).__getattr__(attr)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = RecursiveScriptModule(original_name=LSTM), name = 'forward'

    def __getattr__(self, name: str) -> Union[Tensor, 'Module']:
        if '_parameters' in self.__dict__:
            _parameters = self.__dict__['_parameters']
            if name in _parameters:
                return _parameters[name]
        if '_buffers' in self.__dict__:
            _buffers = self.__dict__['_buffers']
            if name in _buffers:
                return _buffers[name]
        if '_modules' in self.__dict__:
            modules = self.__dict__['_modules']
            if name in modules:
                return modules[name]
        raise ModuleAttributeError("'{}' object has no attribute '{}'".format(
>           type(self).__name__, name))
E       torch.nn.modules.module.ModuleAttributeError: 'RecursiveScriptModule' object has no attribute 'forward'

To Reproduce

import pytest
import torch
from torch import nn
import coremltools as ct

@pytest.fixture
def lstm_model():
    model = nn.LSTM(3, 3)
    model.eval()
    return model

def test_lstm_run(lstm_model):
    seq_len = 12
    inputs = torch.randn(seq_len, 1, 3)
    hidden = (torch.randn(1, 1, 3), torch.randn(1, 1, 3))
    output, hidden = lstm_model(inputs, hidden)
    assert list(output.shape) == [seq_len, 1, 3]

def test_lstm_traced(lstm_model):
    inputs = torch.randn(12, 1, 3)
    traced_model = torch.jit.trace(lstm_model, inputs)

    ct.convert(
        model=traced_model,
        inputs=[
            ct.TensorType(name="sequence", shape=(ct.RangeDim(1, 50), 1, 3))
        ]
    )

def test_lstm_scripted(lstm_model):
    scripted_model = torch.jit.script(lstm_model)

    ct.convert(
        model=scripted_model,
        inputs=[
            ct.TensorType(name="sequence", shape=(ct.RangeDim(1, 50), 1, 3))
        ]
    )

System environment (please complete the following information):

• coremltools version: 4.1
• How you install python (anaconda, virtualenv, system): system
• python version (e.g. 3.7): python3.6
• any other relevant information:
  • torch version: 1.7

[simgt]

Encapsulating the nn.LSTM into a nn.Module subclass actually resolves this first issue. But the process fails further away:

class LstmWrapper(nn.Module):
    def __init__(self, *args):
        super().__init__()
        self.lstm = nn.LSTM(*args)

    def forward(self, inputs):
        return self.lstm(inputs)

Error:

        graph, params = _torch._C._jit_pass_lower_graph(
>           torchscript.forward.graph, torchscript._c
        )
E       RuntimeError: Unknown type bool encountered in graph lowering. This type is not supported in ONNX export.

Related to pytorch/pytorch/issues/41674.

[PareshKamble]

Hi @notsimon Did you find a wayout for this error yet?

[simgt]

Hi @notsimon Did you find a wayout for this error yet?

No I didn't, sorry. Tracing is a good enough alternative for my use case.

[TobyRoseman]

I've made some minor changes to the code in the description of this issue. With these changes, I can reproduce the problem.

As you've mentioned the error here is actually coming from PyTorch when we try to get the expanded and optimize IR.

I'll leave this issue open, but I don't think there is much we can do here. We need the related PyTorch issue (https://github.com/pytorch/pytorch/issues/41674) to be resolved.