转换onnx模型的时候遇到_thnn_fused_lstm_cell的报错,刚开始的报错是提示Keyerror:_thnn_fused_lstm_cell,然后在https://github.com/pytorch/pytorch/issues/25533看到有人说是pytorch和CUDA版本太低的原因以及当使用CPU就不会报错,然后我首先尝试将pytorch升级到最新,CUDA也升级到10.2,最终还是报错Exporting the operator _thnn_fused_lstm_cell to ONNX opset version 9 is not supported. Please open a bug to request ONNX export support for the missing operator.
之后我定位到这个错误的产生是转换onnx模型时使用了torch.nn.LSTMCell()这个方法,如果使用torch.nn.LSTM()就不会报错,但是如果在CPU环境下将torch.nn.LSTMCell()转换成onnx则不会报错。总的来说就是torch.nn.LSTMCell()方法目前不支持在GPU环境下转换成onnx模型。目前已经给pytorch的GitHub提交issue,等待官方出解决方法。
下面是报错复现的列子:
import torch
import torch.nn as nn
class StackedLSTM(nn.Module):
def __init__(self, num_layers, input_size, rnn_size, dropout):
super(StackedLSTM, self).__init__()
self.dropout = nn.Dropout(dropout)
self.num_layers = num_layers
self.layers = nn.ModuleList()
for _ in range(num_layers):
self.layers.append(nn.LSTMCell(input_size, rnn_size))
input_size = rnn_size
def forward(self, input_feed, hidden):
h_0, c_0 = hidden
h_1, c_1 = [], []
for i, layer in enumerate(self.layers):
h_1_i, c_1_i = layer(input_feed, (h_0[i], c_0[i]))
input_feed = h_1_i
if i + 1 != self.num_layers:
input_feed = self.dropout(input_feed)
h_1 += [h_1_i]
c_1 += [c_1_i]
h_1 = torch.stack(h_1)
c_1 = torch.stack(c_1)
return input_feed, (h_1, c_1)
#下面是在CPU环境下转换onnx模型,可以成功转换
lstm = StackedLSTM(2, 580, 500, 0.3) #实例化StackedLSTM
h_t = torch.randn(2, 10, 500) #隐状态h
h_c = torch.randn(2, 10, 500) #隐状态c
hidden = (h_t, h_c)
a = torch.randn(10, 580) #输入
torch.onnx.export(lstm, (a, hidden), f='LSTM.onnx', input_names=['input', 'hidden'], output_names=['output', 'h_t', 'h_c']) #转换onnx模型
下面是在GPU环境下转换onnx模型,报错!
lstm = StackedLSTM(2, 580, 500, 0.3).cuda() #实例化StackedLSTM
h_t = torch.randn(2, 10, 500).cuda() #隐状态h
h_c = torch.randn(2, 10, 500).cuda() #隐状态c
hidden = (h_t, h_c)
a = torch.randn(10, 580).cuda() #输入
torch.onnx.export(lstm, (a, hidden), f='LSTM.onnx', input_names=['input', 'hidden'], output_names=['output', 'h_t', 'h_c']) #转换onnx模型