How to use ONNX to deploy LSTM model? Is RNN supported in ONNX?

Hi, I am working on deploying a pre-trained LSTM model using ONNX. I have obtained the .onnx file following the tutorial of Transfering a model from PyTorch to Caffe2 and Mobile using ONNX.
But for my own model, which is a simple 1-layer LSTM, the error occurs like this:

Traceback (most recent call last):
  File "test.py", line 42, in <module>
    get_onnx_file()
  File "test.py", line 40, in get_onnx_file
    torch_out = torch.onnx.export(model, (x, hc), onnx_filename, verbose=True)
  File "/home/me/anaconda2/lib/python2.7/site-packages/torch/onnx/__init__.py", line 75, in export
    _export(model, args, f, export_params, verbose, training)
  File "/home/me/anaconda2/lib/python2.7/site-packages/torch/onnx/__init__.py", line 102, in _export
    torch._C._jit_pass_onnx(trace)
  File "/home/me/anaconda2/lib/python2.7/site-packages/torch/onnx/__init__.py", line 133, in _run_symbolic_method
    return symbolic_fn(*args)
  File "/home/me/anaconda2/lib/python2.7/site-packages/torch/nn/_functions/rnn.py", line 357, in symbolic
    raise RuntimeError("hack_onnx_rnn NYI")
RuntimeError: hack_onnx_rnn NYI

The code that can reproduce my error is like this:

import torch
import torch.nn as nn
from torch.autograd import Variable
import torch.nn.functional as F

## build a simple LSTM network
class SimpleLSTM(nn.Module):
    def __init__(self, in_size, hidden_size):
        super(SimpleLSTM, self).__init__()
        self.nlayers = 1
        self.nhid = hidden_size
        self.ninp = in_size
        self.rnn = nn.LSTM(in_size, hidden_size, self.nlayers)

    def init_hidden(self, bsz, volatile=False):
        weight = next(self.parameters()).data
        return (Variable(weight.new(self.nlayers, bsz, self.nhid).zero_(), volatile=volatile),
                    Variable(weight.new(self.nlayers, bsz, self.nhid).zero_(), volatile=volatile))
   
    def forward(self, x, hidden):
        out, hidden_tp1 = self.rnn(x, hidden)
        return out, hidden_tp1

in_size, hidden_size = 6, 128

def get_model():
    model = SimpleLSTM(in_size, hidden_size)
    return model

model = get_model()

## get onnx file using `export`
def get_onnx_file():
    onnx_filename = './lstm.onnx'
    # input: x
    x = Variable(torch.rand(1, 1, in_size), volatile=True)
    # input: hidden
    hc = model.init_hidden(1, True)

    torch_out = torch.onnx.export(model, (x, hc), onnx_filename, verbose=True)
# try to output onnx file, where error occurs
get_onnx_file()

Could anybody give me a hint? Thanks!

PS: I found this in documentation page of torch.onnx: (Here)

We plan on expanding support to more operators; RNNs are high on our priority list.

So RNN(i.e. RNN, LSTM/GRU) has not been supported by ONNX yet?

Got the same problem here. On the docs, v0.3 should’ve already supported RNN operation, but I still get the same error. Do we need to install from master?

I think you can have a try. I gave up using ONNX at that moment. I exported the weight of the LSTM model to MXNet and used tvm.

How did you export the weight to MXNet? Did you re-trained the net with MXNet directly?

I didn’t retrain the model using MXNet. Just fetched the weights from the LSTM model and filled them into MXNet params. It is straightforward.

I see, ok then. I’ll try it in my case. Thanks!

You’d better first try to re-build pytorch by the source code of master branch. Glad to help you.

Getting this same error. Anybody know if RNN’s are officially supported yet?

1 Like