如何理解Keras预训练权重文件？加载到PyTorch遇报错

错误原因 #

torch.load仅支持加载PyTorch自身序列化的权重文件（如.pth/.pt格式），而TensorFlow的.h5是HDF5格式存储文件，两者序列化机制完全不同，直接调用torch.load必然触发UnpicklingError。

解决步骤 #

1. 理解.h5文件的结构 #

你看到的add_loss、dense等HDF5组，对应TensorFlow模型的各个层或模块，每组下的数据集就是该层的权重参数（比如dense/kernel:0是全连接层的权重矩阵，dense/bias:0是偏置项）。

2. 手动映射权重到PyTorch模型 #

核心是保证PyTorch模型的结构与TensorFlow模型完全对齐，再将对应层的权重做形状转换后赋值（两者部分层的参数维度顺序存在差异）。

示例代码（基于h5py读取）

import h5py
import torch
import torch.nn as nn

# 读取HDF5权重文件
with h5py.File('pretrained.h5', 'r') as h5_f:
    # 定义与TensorFlow结构一致的PyTorch模型
    class TF2TorchModel(nn.Module):
        def __init__(self):
            super().__init__()
            self.dense1 = nn.Linear(784, 256)  # 对应TF的dense层
            self.relu = nn.ReLU()
            self.dense2 = nn.Linear(256, 10)   # 对应TF的dense_1层
        
        def forward(self, x):
            x = self.relu(self.dense1(x))
            return self.dense2(x)

    model = TF2TorchModel()

    # 映射全连接层权重：TF权重形状是[in_features, out_features]，PyTorch是[out_features, in_features]，需转置
    # 处理dense1层
    tf_dense1_kernel = h5_f['dense/kernel:0'][()]
    tf_dense1_bias = h5_f['dense/bias:0'][()]
    model.dense1.weight.data = torch.tensor(tf_dense1_kernel.T, dtype=torch.float32)
    model.dense1.bias.data = torch.tensor(tf_dense1_bias, dtype=torch.float32)

    # 处理dense2层
    tf_dense2_kernel = h5_f['dense_1/kernel:0'][()]
    tf_dense2_bias = h5_f['dense_1/bias:0'][()]
    model.dense2.weight.data = torch.tensor(tf_dense2_kernel.T, dtype=torch.float32)
    model.dense2.bias.data = torch.tensor(tf_dense2_bias, dtype=torch.float32)

# 验证权重加载
print(model.dense1.weight.shape)  # 应输出torch.Size([256, 784])

卷积层的特殊处理

如果模型包含卷积层，注意两者的卷积核维度顺序差异：

• TensorFlow卷积核形状：[height, width, in_channels, out_channels]
• PyTorch卷积核形状：[out_channels, in_channels, height, width]
  转换方式示例：

# 假设TF卷积层权重存在conv2d/kernel:0
tf_conv_kernel = h5_f['conv2d/kernel:0'][()]
# 转置维度适配PyTorch
torch_conv_kernel = tf_conv_kernel.transpose(3, 2, 0, 1)
model.conv1.weight.data = torch.tensor(torch_conv_kernel, dtype=torch.float32)

3. 简便方式：用TensorFlow加载后提取权重 #

如果是Keras保存的.h5模型，可以先用TensorFlow加载完整模型，再提取权重列表按顺序映射：

import tensorflow as tf
import torch

# TensorFlow加载模型并提取权重
tf_model = tf.keras.models.load_model('pretrained.h5')
tf_weights_list = tf_model.get_weights()

# 遍历PyTorch模型参数，按顺序赋值（注意形状转换）
torch_params = list(model.parameters())
for idx, param in enumerate(torch_params):
    tf_weight = tf_weights_list[idx]
    # 针对全连接/卷积层做形状转换，其他层直接赋值
    if len(tf_weight.shape) == 2:  # 全连接层权重
        param.data = torch.tensor(tf_weight.T, dtype=torch.float32)
    elif len(tf_weight.shape) == 4:  # 卷积层权重
        param.data = torch.tensor(tf_weight.transpose(3,2,0,1), dtype=torch.float32)
    else:  # 偏置、BN的均值/方差等，形状一致直接赋值
        param.data = torch.tensor(tf_weight, dtype=torch.float32)

关键注意事项 #

• 必须保证PyTorch模型的结构（层类型、输入输出维度、激活函数、连接方式）与TensorFlow模型完全一致，否则权重映射会失效。
• 注意数据类型匹配：TensorFlow权重默认是float32，PyTorch也要用对应类型，避免精度问题。
• 批量归一化层的参数可直接映射：TF的gamma对应PyTorch的weight，beta对应bias，moving_mean对应running_mean，moving_var对应running_var。

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