PyTorch LSTM状态参数向Keras LSTM转换方法求助

我之前刚好踩过这个参数转换的坑，咱们先把两者的参数差异说透，再给你具体的转换步骤和代码。

首先得明确：PyTorch和Keras的LSTM核心计算逻辑是一致的，但参数拆分方式不同：

• PyTorch把LSTM的偏置拆成了输入侧偏置（bias_ih）和隐藏状态侧偏置（bias_hh），加上输入权重（weight_ih）、隐藏权重（weight_hh），一共4组参数；
• Keras则把输入侧和隐藏侧的偏置合并成了一组偏置（bias），加上输入权重（kernel）、隐藏权重（recurrent_kernel），一共3组参数。

另外还要注意权重的维度顺序：PyTorch的权重矩阵是(4*hidden_size, input_size/hidden_size)，而Keras是(input_size/hidden_size, 4*hidden_size)，本质是因为两者的线性计算顺序不同（PyTorch是x @ W.T，Keras是W @ x），所以需要转置权重。

针对你的双向LSTM案例的转换步骤 #

你的模型参数：hidden_size=64（所以4*64=256）、input_size=512，双向结构包含forward和backward两个分支，我们分别处理：

1. Forward分支参数转换

从PyTorch提取forward分支的4组参数：

• rnn.0.rnn.weight_ih_l0（shape: (256,512)）→ Keras的bidirectional_1/forward_lstm_1/kernel：转置这个矩阵，得到shape (512,256)
• rnn.0.rnn.weight_hh_l0（shape: (256,64)）→ Keras的bidirectional_1/forward_lstm_1/recurrent_kernel：转置这个矩阵，得到shape (64,256)
• rnn.0.rnn.bias_ih_l0 + rnn.0.rnn.bias_hh_l0 → Keras的bidirectional_1/forward_lstm_1/bias：把两个(256,)的向量逐元素相加，得到shape (256,)

2. Backward分支参数转换

同理处理backward分支：

• rnn.0.rnn.weight_ih_l0_reverse（shape: (256,512)）→ Keras的bidirectional_1/backward_lstm_1/kernel：转置得到(512,256)
• rnn.0.rnn.weight_hh_l0_reverse（shape: (256,64)）→ Keras的bidirectional_1/backward_lstm_1/recurrent_kernel：转置得到(64,256)
• rnn.0.rnn.bias_ih_l0_reverse + rnn.0.rnn.bias_hh_l0_reverse → Keras的bidirectional_1/backward_lstm_1/bias：逐元素相加得到(256,)

代码示例（PyTorch → Keras 参数迁移） #

假设你已经加载了PyTorch模型和定义好了对应的Keras模型，代码大概是这样：

import torch
from tensorflow import keras

# 加载PyTorch预训练模型
pytorch_model = torch.load("your_pytorch_model.pth")
# 定义好对应的Keras双向LSTM模型（结构要完全匹配：input_size=512, hidden_size=64, bidirectional=True）
keras_model = your_defined_keras_model()

# 提取PyTorch的forward分支参数
weight_ih_forward = pytorch_model.state_dict()['rnn.0.rnn.weight_ih_l0'].numpy()
weight_hh_forward = pytorch_model.state_dict()['rnn.0.rnn.weight_hh_l0'].numpy()
bias_ih_forward = pytorch_model.state_dict()['rnn.0.rnn.bias_ih_l0'].numpy()
bias_hh_forward = pytorch_model.state_dict()['rnn.0.rnn.bias_hh_l0'].numpy()

# 转换forward分支参数给Keras
keras_model.get_layer('bidirectional_1').forward_layer.set_weights([
    weight_ih_forward.T,  # 转置权重
    weight_hh_forward.T,  # 转置权重
    bias_ih_forward + bias_hh_forward  # 合并偏置
])

# 提取并转换backward分支参数
weight_ih_backward = pytorch_model.state_dict()['rnn.0.rnn.weight_ih_l0_reverse'].numpy()
weight_hh_backward = pytorch_model.state_dict()['rnn.0.rnn.weight_hh_l0_reverse'].numpy()
bias_ih_backward = pytorch_model.state_dict()['rnn.0.rnn.bias_ih_l0_reverse'].numpy()
bias_hh_backward = pytorch_model.state_dict()['rnn.0.rnn.bias_hh_l0_reverse'].numpy()

keras_model.get_layer('bidirectional_1').backward_layer.set_weights([
    weight_ih_backward.T,
    weight_hh_backward.T,
    bias_ih_backward + bias_hh_backward
])

# 保存转换后的Keras模型
keras_model.save("converted_keras_model.h5")

验证转换正确性 #

转换后建议用相同的输入数据分别喂给PyTorch模型和Keras模型，检查输出的差异（应该非常小，因为浮点精度问题可能有微小误差），确保参数迁移正确。

内容的提问来源于stack exchange，提问作者harish2704