首页
学习
活动
专区
工具
TVP
发布
社区首页 >问答首页 >如何正确地将pytorch LSTM转换为keras CuDNNLSTM?

如何正确地将pytorch LSTM转换为keras CuDNNLSTM?
EN

Stack Overflow用户
提问于 2019-05-21 11:00:46
回答 1查看 821关注 0票数 0

我正在尝试手动将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无论如何都使用相同的后端。

EN

回答 1

Stack Overflow用户

回答已采纳

发布于 2019-05-22 03:14:28

更新:解决方案是将torch模型转换为keras基础LSTM模型,然后调用

base_lstm_model.save_weights('1.h5')
cudnn_lstm_model.load_weights('1.h5')
票数 1
EN
页面原文内容由Stack Overflow提供。腾讯云小微IT领域专用引擎提供翻译支持
原文链接:

https://stackoverflow.com/questions/56230511

复制
相关文章

相似问题

领券
问题归档专栏文章快讯文章归档关键词归档开发者手册归档开发者手册 Section 归档