PyTorch保存的.h5模型无法被TensorFlow加载的问题求助

错误原因 #

PyTorch的torch.save()不管是保存state_dict还是整个模型，生成的都是PyTorch专属的二进制序列化文件，哪怕你把后缀改成.h5，本质也不是TensorFlow支持的HDF5格式模型文件。h5py加载时会校验文件签名，自然会抛出OSError: Unable to open file (file signature not found)错误。

正确的跨框架转换方案 #

方案1：通过ONNX中间格式转换 #

这是跨框架模型迁移的标准方式，适配大多数常见模型：

步骤1：PyTorch模型导出为ONNX

import torch

# 模型切换到评估模式
model.eval()
# 构造匹配模型输入维度的示例张量（根据你的模型输入调整，比如这里是batch=1，3通道，224x224图像）
dummy_input = torch.randn(1, 3, 224, 224)
# 导出ONNX模型，opset_version选兼容的版本（推荐11/13）
torch.onnx.export(model, dummy_input, "model.onnx", opset_version=11)

步骤2：ONNX转TensorFlow的.h5格式

import onnx
from onnx_tf.backend import prepare

# 加载ONNX模型
onnx_model = onnx.load("model.onnx")
# 转换为TensorFlow可执行模型
tf_rep = prepare(onnx_model)
# 导出为.h5格式
tf_rep.export_graph("model.h5")

方案2：手动迁移权重（适合简单模型） #

如果模型结构简单，可以手动提取PyTorch权重并赋值到TensorFlow模型中，注意框架间的参数维度差异：

import torch
import tensorflow as tf

# 1. PyTorch端提取权重
pt_weights = model.state_dict()

# 2. TensorFlow端构建与PyTorch结构一致的模型tf_model
# 示例：假设模型是ResNet50，需提前构建结构
# tf_model = tf.keras.applications.ResNet50(weights=None, input_shape=(224,224,3))

# 3. 逐层赋值权重（注意调整维度顺序，比如卷积层）
for layer_name, param in pt_weights.items():
    # 跳过不需要手动赋值的参数（如BatchNorm的running_mean可按需处理）
    if "weight" in layer_name:
        # PyTorch卷积权重维度：(out_channels, in_channels, h, w)
        # TensorFlow卷积权重维度：(h, w, in_channels, out_channels)
        adjusted_param = param.numpy().transpose(2, 3, 1, 0)
        # 适配TensorFlow的层命名规则（替换PyTorch的.为/）
        tf_model.get_layer(layer_name.replace(".", "/")).set_weights([adjusted_param])
    elif "bias" in layer_name:
        adjusted_param = param.numpy()
        tf_model.get_layer(layer_name.replace(".", "/")).set_weights([adjusted_param])

# 4. 保存TensorFlow模型
tf_model.save("model.h5")

注意事项 #

• ONNX转换时，opset版本要和PyTorch版本匹配，高版本PyTorch建议用更高的opset（比如13+）。
• 自定义层在ONNX转换时可能报错，需要先将自定义层替换为ONNX支持的标准算子。
• 手动迁移权重时，要仔细核对各层的参数维度和命名，不同框架的层命名规则可能有差异。

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