pytorch模型转onnxExportingtheoperator_thnn_fused_lstm_celltoONNXopsetversion9isnotsupported

转换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 nnclass 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_layersself.layers = nn.ModuleList()for _ in range(num_layers):self.layers.append(nn.LSTMCell(input_size, rnn_size))input_size = rnn_sizedef forward(self, input_feed, hidden):h_0, c_0 = hiddenh_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_iif 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=&＃39;LSTM.onnx&＃39;, input_names=[&＃39;input&＃39;, &＃39;hidden&＃39;], output_names=[&＃39;output&＃39;, &＃39;h_t&＃39;, &＃39;h_c&＃39;]) #转换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=&＃39;LSTM.onnx&＃39;, input_names=[&＃39;input&＃39;, &＃39;hidden&＃39;], output_names=[&＃39;output&＃39;, &＃39;h_t&＃39;, &＃39;h_c&＃39;]) #转换onnx模型