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?