如何正确地将pytorch LSTM转换为keras cudnlstm？

Question

我正在尝试手动将Pythorch模型转换为Tensorflow进行部署。ONNX似乎并没有从Pytorch lstm到Tensorflow cudnlstms，所以我要手工编写它。在

我试过下面的代码： 这是在一个运行python2.7、Pytorch 1.0、tensorflow 1.12、cuda9的Python环境中运行的。我在Pytorch层中运行它时没有任何偏差，因为它遵循一个batchnorm，但是由于Keras没有提供这个选项，所以我只是指定了一个0偏差。在

import torch
import tensorflow as tf
import numpy as np
from tensorflow.keras.layers import CuDNNLSTM, Bidirectional
from tensorflow.keras.models import Sequential, Model

input_size = 80
hidden_size = 512
with torch.no_grad():
    rnn1 = torch.nn.LSTM(input_size=input_size, hidden_size=hidden_size, bidirectional=True, bias=False, batch_first=True).cuda()

model = Sequential()
model.add(Bidirectional(CuDNNLSTM(hidden_size, return_sequences=True),  input_shape=(None, input_size), name='rnn'))

bias_size = rnn1.weight_hh_l0.detach().cpu().numpy().T.shape[1] * 2
keras_format_weights = [
                    rnn1.weight_ih_l0.detach().cpu().numpy().T,
                    rnn1.weight_hh_l0.detach().cpu().numpy().T,
                    np.zeros(bias_size,),
                    rnn1.weight_ih_l0_reverse.detach().cpu().numpy().T,
                    rnn1.weight_hh_l0_reverse.detach().cpu().numpy().T,
                    np.zeros(bias_size,),
                  ]


model.layers[0].set_weights(keras_format_weights)

random_test = np.random.rand(1, 1, 80)

res1, _ = rnn1.forward(torch.FloatTensor(random_test).cuda())
res1 = res1.detach().cpu().numpy()
res2 = model.predict(random_test)

print(np.allclose(res1, res2, atol=1e-2))
print(res1)
print(res2)

^{pr2}$

现在，这确实适用于通用的Keras LSTM：

model = Sequential()
model.add(Bidirectional(LSTM(hidden_size, recurrent_activation='sigmoid', return_sequences=True),  input_shape=(None, input_size), name='rnn'))

bias_size = rnn1.weight_hh_l0.detach().cpu().numpy().T.shape[1]
keras_format_weights = [
                    rnn1.weight_ih_l0.detach().cpu().numpy().T,
                    rnn1.weight_hh_l0.detach().cpu().numpy().T,
                    np.zeros((bias_size,)),
                    rnn1.weight_ih_l0_reverse.detach().cpu().numpy().T,
                    rnn1.weight_hh_l0_reverse.detach().cpu().numpy().T,
                    np.zeros((bias_size,))
                  ]

但我需要CuDNNLSTM的速度优势，而且Pytorch也在使用相同的后端。在

Answer 1

更新：解决方案是将torch模型转换为keras-baselstm模型，然后调用

base_lstm_model.save_weights('1.h5')
cudnn_lstm_model.load_weights('1.h5')