Skip to main content
 首页 » 编程设计

python之手动将 LSTM 从 Tensorflow 导入 PyTorch

2024年11月24日30zengkefu

我正在尝试将预训练模型从 tensorflow 导入 PyTorch。它采用单个输入并将其映射到单个输出。 当我尝试导入 LSTM 权重时,出现了困惑

我使用以下函数从文件中读取权重及其变量:

def load_tf_model_weights():         
 
    modelpath = 'models/model1.ckpt.meta' 
 
    with tf.Session() as sess:         
        tf.train.import_meta_graph(modelpath)  
        init = tf.global_variables_initializer() 
        sess.run(init)   
        vars = tf.trainable_variables()         
        W = sess.run(vars) 
 
    return W,vars 
 
W,V = load_tf_model_weights() 

然后我正在检查重物的形状

In [33]:  [w.shape for w in W] 
Out[33]: [(51, 200), (200,), (100, 200), (200,), (50, 1), (1,)] 

此外变量定义为

In [34]:    V 
Out[34]:  
[<tf.Variable 'rnn/multi_rnn_cell/cell_0/lstm_cell/kernel:0' shape=(51, 200) dtype=float32_ref>, 
<tf.Variable 'rnn/multi_rnn_cell/cell_0/lstm_cell/bias:0' shape=(200,) dtype=float32_ref>, 
<tf.Variable 'rnn/multi_rnn_cell/cell_1/lstm_cell/kernel:0' shape=(100, 200) dtype=float32_ref>, 
<tf.Variable 'rnn/multi_rnn_cell/cell_1/lstm_cell/bias:0' shape=(200,) dtype=float32_ref>, 
<tf.Variable 'weight:0' shape=(50, 1) dtype=float32_ref>, 
<tf.Variable 'FCLayer/Variable:0' shape=(1,) dtype=float32_ref>] 

所以我可以说 W 的第一个元素定义了 LSTM 的内核,第二个元素定义了它的偏差。根据this post ,内核的形状定义为 [input_depth + h_depth, 4 * self._num_units] 偏差为 [4 * self._num_units]。我们已经知道 input_depth 是 1。所以我们得到,h_depth 和 _num_units 都有值 50。

在 pytorch 中,我想为其分配权重的 LSTMCell 如下所示:

In [38]: cell = nn.LSTMCell(1,50) 
In [39]: [p.shape for p in cell.parameters()] 
Out[39]:  
[torch.Size([200, 1]), 
torch.Size([200, 50]), 
torch.Size([200]), 
torch.Size([200])] 

前两个条目可以被形状为 (51,200) 的 W 的第一个值覆盖。但是来自 Tensorflow 的 LSTMCell 只产生一个形状为 (200) 的偏置,而 pytorch 需要两个偏置

通过排除偏差,我还剩下权重:

cell2 = nn.LSTMCell(1,50,bias=False) 
[p.shape for p in cell2.parameters()] 
Out[43]: [torch.Size([200, 1]), torch.Size([200, 50])] 

谢谢!

请您参考如下方法:

pytorch 使用 CuDNN 的 LSTM 底层(即使您没有 CUDA,它仍然使用兼容的东西)因此它有一个额外的偏置项。

因此,您可以选择两个总和等于 1 的数字(0 和 1、1/2 和 1/2 或其他任何值),并将您的 pytorch 偏差设置为这些数字乘以 TF 的偏差。

pytorch_bias_1 = torch.from_numpy(alpha * tf_bias_data) 
pytorch_bias_2 = torch.from_numpy((1.0-alpha) * tf_bias_data)