Pytorch模型转Tensorflow模型

引言

最近收到领导布置的一个小任务,需要将pth文件转换到tensorflow模型。这里将采用领导给的pytorchocr代码,来记录下如何进行模型转换。

一、基本思想

想要将pytorch训练得到的pth文件转换到tensorflow的pb文件,基本思想就是先将pth转成onnx,再又onnx转换到pb文件。整体没什么难度,我在下面直接给出代码。

二、转换代码

转换代码如下,我这里采用的是pytorchocr网络,有了具体网络结构就可以进行onnx转换,最后到tensorflow转换。

复制代码
import torch
import torchvision
import tensorflow as tf
from onnx_tf.backend import prepare
from pytorchocr.base_ocr_v20 import BaseOCRV20
import onnx

det_model_path = "E:\File\ckpt\pytorch\det.pth"
# det_model_path = r"D:\File_save\pycharm\demo\Script\ckpt\pytorch\rec.pth"

class PPOCRv3DetConverter(BaseOCRV20):
    def __init__(self, config, **kwargs):
        super(PPOCRv3DetConverter, self).__init__(config, **kwargs)
        # self.load_paddle_weights(paddle_pretrained_model_path)
        self.net.eval()
        # self.training = False


cfg = {'model_type':'det',
           'algorithm':'DB',
           'Transform':None,
           'Backbone':{'name':'MobileNetV3', 'model_name':'large', 'scale':0.5, 'disable_se':True},
           'Neck':{'name':'RSEFPN', 'out_channels':96, 'shortcut': True},
           'Head':{'name':'DBHead', 'k':50}

       }

# cfg = {'model_type': 'rec',
#        'algorithm': 'CRNN',
#        'Transform': None,
#        'Backbone': {'name': 'MobileNetV1Enhance',
#                     'scale': 0.5,
#                     'last_conv_stride': [1, 2],
#                     'last_pool_type': 'avg'},
#        'Neck': {'name': 'SequenceEncoder',
#                 'dims': 64,
#                 'depth': 2,
#                 'hidden_dims': 120,
#                 'use_guide': True,
#                 'encoder_type': 'svtr'},
#        'Head': {'name': 'CTCHead', 'fc_decay': 2e-05}
#        }

converter = PPOCRv3DetConverter(cfg)
converter.load_state_dict(torch.load(det_model_path))
# det
dummy_input = torch.randn(1, 3, 640, 640)
# rec
# dummy_input = torch.randn(1, 3, 48, 320)

det_model = converter.net

# 设置模型为评估模式
converter.net.eval()

# 导出模型为ONNX格式
# torch.onnx.export(det_model, dummy_input, 'rec_model.onnx', verbose=False, opset_version=11)
# torch.onnx.export(det_model, dummy_input, 'det_model.onnx', verbose=False, opset_version=12)

# 加载ONNX模型
onnx_model = onnx.load('E:\File\ckpt\onnx\det_model.onnx')
# onnx_model = onnx.load(r'D:\File_save\pycharm\demo\Script\ckpt\rec_model.onnx')
#
# 转换为TensorFlow模型
tf_model = prepare(onnx_model, strict=False)
#
# 保存为TensorFlow模型文件
tf_model.export_graph('det_model.pb')
# tf_model.export_graph('rec_model.pb')
# #
# tf.compat.v1.reset_default_graph()
# with tf.compat.v1.Session() as sess:
#     with tf.gfile.GFile('det_model.onnx', 'rb') as f:
#         graph_def = tf.compat.v1.GraphDef()
#         graph_def.ParseFromString(f.read())
#         tf.import_graph_def(graph_def, name='')
#
#     # 保存为.pb文件
#     tf.io.write_graph(sess.graph, '', 'det_model.pb', as_text=False)
相关推荐
极限实验室1 小时前
Coco AI 实战(一):Coco Server Linux 平台部署
人工智能
杨过过儿1 小时前
【学习笔记】4.1 什么是 LLM
人工智能
巴伦是只猫1 小时前
【机器学习笔记Ⅰ】13 正则化代价函数
人工智能·笔记·机器学习
大千AI助手1 小时前
DTW模版匹配:弹性对齐的时间序列相似度度量算法
人工智能·算法·机器学习·数据挖掘·模版匹配·dtw模版匹配
AI生存日记2 小时前
百度文心大模型 4.5 系列全面开源 英特尔同步支持端侧部署
人工智能·百度·开源·open ai大模型
LCG元2 小时前
自动驾驶感知模块的多模态数据融合:时序同步与空间对齐的框架解析
人工智能·机器学习·自动驾驶
why技术2 小时前
Stack Overflow,轰然倒下!
前端·人工智能·后端
烛阴3 小时前
简单入门Python装饰器
前端·python
超龄超能程序猿3 小时前
(三)PS识别:基于噪声分析PS识别的技术实现
图像处理·人工智能·计算机视觉