PyTorch模型转换为TensorFlow模型的教程

前言

目前大多数模型都是pytorch格式，部署上很多tfserving用的比较多，因此模型格式需要是save_model.pb的格式，本篇文章介绍将pytorch转化为tensorflow格式模型方式。

核心过程：pytorch====>onnx====>tensorflow

一、pytorch转onnx

第一步要先加载模型，初始化模型网络，加载模型权重之类的。比如这里的我要加载instructor-large模型文件。instructor-large是一个文件夹，我们来看一下这里面的内容。

模型网络是T5EncorderModel，加载代码如下所示。

登录后复制

from transformers import T5EncoderModel
import torch
model = T5EncoderModel.from_pretrained("../instructor-large")

加载顺利以后，使用torch.onnx.export函数转换。

登录后复制

torch.onnx.export(model, torch.zeros(1, 512, dtype=torch.long), "./model.onnx", 
input_names=["input"], output_names=["output"], opset_version=12)

核心前三个参数：第一个是model，第二个是model的输入格式，第三个是目标路径
还要再加一个： opset_version=12，一定要加这个，不然后面会报错，去官网查了是因为这是一个官方的bug，也不能说是bug吧，就是版本之间的问题，opset_version设置为12。

执行好上面的export代码以后，就会生成onnx模型（这里指挥生成model.onnx一个文件，其他的不用管）

到这里pytorch就成功转化为onnx了。（注意模型大小不能超过2GB，超过了转化会不成功，是因为官方还不支持超过2GB大小的模型）

二、onnx转tensorflow

首先是读取onnx模型

登录后复制

import onnx
import tf2onnx
from onnx_tf.backend import prepare
# 加载ONNX模型
onnx_model = onnx.load('model.onnx', load_external_data=False)

然后执行以下代码即可

登录后复制

tf_pb = prepare(onnx_model)
tf_pb.export_graph('model_tensorflow')

注意这里如何前面不设置opset_version为12的话会出现以下错误，依然是官方的问题哈哈哈。

最终如果你成功执行完上面的代码就会生成一个model_tensorflow文件夹，里面长这样。就大功告成了。

免责声明：本文系网络转载或改编，未找到原创作者，版权归原作者所有。如涉及版权，请联系删