引言
最近收到领导布置的一个小任务,需要将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)