Как правильно конвертировать pytorch LSTM в керас CuDNNLSTM?

Question

Я пытаюсь вручную преобразовать модель Pytorch в Tensorflow для развертывания. Похоже, что ONNX изначально не переходил от LSTM Pytorch к Tensorflow CuDNNLSTM, поэтому я пишу это вручную.

Я попробовал код ниже: Это выполняется в среде Anaconda, работающей на Python 2.7, Pytorch 1.0, тензор потока 1.12, cuda9. Я выполняю это без смещения в слое Pytorch, так как он следует за пакетной формой, но поскольку 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)

False
[[[ 0.01265562  0.07478553  0.0470101  ... -0.02260824  0.0243004
   -0.0261014 ]]]
[[[-0.05316251 -0.00230848  0.03070898 ...  0.01497027  0.00976444
   -0.01095549]]]

Теперь, это работает с универсальным 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

Обновление: решение состояло в том, чтобы преобразовать модель факела в базовую модель LSTM keras, а затем вызвать

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