Pytorch模型转Tensorflow模型的那些事

最新推荐文章于 2024-08-07 09:09:41 发布

AmibitionWei

最新推荐文章于 2024-08-07 09:09:41 发布

阅读量1w

点赞数 14

分类专栏：励志专栏深度学习文章标签： tensorflow 深度学习 pytorch 神经网络

本文链接：https://blog.csdn.net/qq_30666517/article/details/113249029

版权

本文介绍如何将PyTorch模型转换为TensorFlow模型进行部署。通过两个步骤实现：首先是将PyTorch模型文件(.pt)转换为ONNX文件，接着将ONNX文件转换为TensorFlow的.pb文件。文章提供了代码示例，并讨论了版本兼容性和环境配置等问题。

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

Pytorch模型转Tensorflow模型部署上线

pytorch模型转tensorflow流程

最近在研究Query2Title模型，学术界上快速实验一般都用pytorch，但是业界部署模型上大多都还是tensorflow模型部署。也可能是自己太懒了，哈哈，pytorch用久了tensorflow有点生疏，懒的去用啦！结果就…，线上部署必须得tensorflow模型，所以就把torch的.pt文件，转成tensorflow的pb文件把，中间遇到一些问题记录下来，仅供借鉴参考！<ps!立下flag，坚持写一些博客>

pytorch模型转tensorflow流程

总体上pytorch转tensorflow遵循2个步骤：

torch模型文件.pt保存成onnx模型文件.onnx
.onnx模型文件转换成tensorflow .pb文件

torch模型文件转onnx文件

这里直接利用torch的onnx库将.pt文件load后导出为.onnx文件即可，基本上没有什么坑，如下：

import tensorflow as tf
import torch
import onnx
from onnx_tf.backend import prepare
import os
import numpy as np


def torch2Onnx():
    """
    pytorch转onnx
    """
    model_pytorch = YourModel()
    model_pytorch.load_state_dict(torch.load('xxx.pt',
                                       map_location='cpu'))
    # 输入placeholder
    dummy_input = torch.randint(0, 10000, (1, 20))
    dummy_output = model_pytorch(dummy_input)
    print(dummy_output.shape)

    # Export to ONNX format
    torch.onnx.export(model_pytorch, 
                      dummy_input, 
                      'model.onnx', 
                      input_names=['inputs'], 
                      output_name

最低0.47元/天解锁文章