pytorch模型转onnx Exporting the operator _thnn_fused_lstm_cell to ONNX opset version 9 is not supported

转换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模型