我正在尝试手动将Pytorch模型转换为Tensorflow进行部署。ONNX似乎不是天生就能从Pytorch LSTM到Tensorflow CuDNNLSTM的,所以这就是我手写它的原因。
我尝试过下面的代码:这是在运行Python2.7,Pytorch1.0,TensorFlow1.12,cuda9的Anaconda环境中运行的。我在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无论如何都使用相同的后端。
发布于 2019-05-22 03:14:28
更新:解决方案是将torch模型转换为keras基础LSTM模型,然后调用
base_lstm_model.save_weights('1.h5')
cudnn_lstm_model.load_weights('1.h5')
https://stackoverflow.com/questions/56230511
复制相似问题