torch模型转tensorflow模型汇总

最新推荐文章于 2025-03-11 14:30:19 发布

Hangfei_Liu

最新推荐文章于 2025-03-11 14:30:19 发布

阅读量2.2k

点赞数 1

分类专栏：模型转换文章标签：深度学习 pytorch tensorflow 人工智能

本文链接：https://blog.youkuaiyun.com/qq_40637261/article/details/120548068

版权

本文详细介绍了如何将训练好的PyTorch模型转换为TensorFlow模型，包括使用onnx-tf和pytorch2keras模块的方法，以及在训练过程中两个框架之间的相互调用。还提到了转换过程中需要注意的细节，如数据归一化、权重维度、BN层的处理等，以减少转换带来的偏差。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

torch模型转tensorflow模型汇总

1，针对已训练的模型，有两种思路：
(1) 通过onnx-tf模块转换：
a，具体流程：
pth=>onnx=>pb
b，环境配置：
conda create -n pth_pb python=3.7
pip install tensorflow2.1.0 -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install tensorflow-addons0.9.1 -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install onnx-tf1.5.0 -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install onnx1.6.0 -i https://pypi.tuna.tsinghua.edu.cn/simple
conda install pytorch torchvision #(这里我用的版本为pytorch1.4, torchvision0.5.0)
c，代码：
pth转onnx

import torchvision
import torch.onnx
import torch.nn as nn
def resnet50():
    model = torchvision.models.resnet50(pretrained=False)
    model.fc = nn.Linear(2048, 2)
    return model
model = resnet50()
model = torch.nn.DataParallel(model)
pthfile = my.pth'
loaded_model = torch.load(pthfile, map_location='cpu')
model.load_state_dict(loaded_model['state_dict'])
input = torch.randn(1, 3, 200, 200)
input_names = ["head_input"]
output_names = ["output"]
onnx_filename = “my.onnx"
torch.onnx.export(model.module, input, onnx_filename, verbose=True, input_names=input_names, output_names=output_names)