如何提取PyTorch预训练模型权重并迁移到TensorFlow中使用

步骤1：准备运行环境 #

先安装需要的依赖库：

• PyTorch：用于加载预训练模型、提取权重
• TensorFlow：用于构建目标模型、载入转换后的权重
• NumPy：用于权重维度适配处理
  安装命令参考：

pip install torch tensorflow numpy

步骤2：加载PyTorch预训练模型并提取权重 #

运行指定代码加载目标LSTM模型，导出所有权重参数：

import torch
import numpy as np

# 加载指定预训练LSTM模型
pytorch_lstm = torch.hub.load("BruceWen120/medal", "lstm")
# 切换到评估模式，避免dropout等训练层影响权重
pytorch_lstm.eval()

# 提取所有参数转成numpy格式存储
pytorch_weights = {}
for name, param in pytorch_lstm.named_parameters():
    pytorch_weights[name] = param.detach().cpu().numpy()

你可以先打印pytorch_weights.keys()查看所有权重的命名、形状，确认模型的层数、隐层维度、是否双向、是否带偏置等核心参数，方便后续构建匹配的TensorFlow模型。

步骤3：构建结构完全对齐的TensorFlow LSTM模型 #

必须保证TensorFlow侧的LSTM结构和PyTorch侧完全一致，否则权重无法正常载入。示例如下（请根据你上一步查到的实际参数调整）：

import tensorflow as tf

# 示例为单层单向LSTM，隐层维度256、带偏置、batch_first=True，按需调整参数
tf_lstm = tf.keras.layers.LSTM(
    units=256,
    return_sequences=True,
    return_state=True,
    use_bias=True,
    batch_input_shape=(None, None, 实际输入特征维度)
)
# 先跑一次样例输入完成模型初始化，给权重分配存储空间
sample_input = tf.random.normal((1, 10, 实际输入特征维度))
_ = tf_lstm(sample_input)

步骤4：权重维度适配与载入 #

PyTorch和TensorFlow的LSTM权重排布逻辑不同，需要做维度转换：

1. PyTorch的LSTM输入权重、隐层权重形状为(隐层大小*4, 输入维度)、(隐层大小*4, 隐层大小)，偏置分为输入偏置、隐层偏置两个，形状都是(隐层大小*4,)
2. TensorFlow的LSTM输入权重、隐层权重形状为(输入维度, 隐层大小*4)、(隐层大小, 隐层大小*4)，偏置合并为一个，形状为(隐层大小*4,)
   转换示例：

# 按你实际的权重名替换下面的参数名
ih_weight = pytorch_weights['lstm.weight_ih_l0'].T # 转置适配TensorFlow维度
hh_weight = pytorch_weights['lstm.weight_hh_l0'].T # 转置适配TensorFlow维度
# 合并两个偏置
bias = pytorch_weights['lstm.bias_ih_l0'] + pytorch_weights['lstm.bias_hh_l0']

# 给TensorFlow模型赋值权重
tf_lstm.set_weights([ih_weight, hh_weight, bias])

如果是多层LSTM或者双向LSTM，按相同逻辑处理每一层、每一个方向的权重即可。

步骤5：验证转换正确性 #

用相同的输入分别喂给两个模型，对比输出误差，误差在1e-5以内就说明转换正常：

# 构造测试输入
test_input = np.random.randn(2, 10, 实际输入特征维度).astype(np.float32)

# PyTorch侧推理结果
with torch.no_grad():
    pt_output, _ = pytorch_lstm(torch.from_numpy(test_input))
pt_output = pt_output.numpy()

# TensorFlow侧推理结果
tf_output, _ = tf_lstm(tf.convert_to_tensor(test_input))
tf_output = tf_output.numpy()

# 打印误差
print("平均误差：", np.mean(np.abs(pt_output - tf_output)))

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