产品
最近在研究Query2Title模型，学术界上快速实验一般都用pytorch，但是业界部署模型上大多都还是tensorflow模型部署。也可能是自己太懒了，哈哈，pytorch用久了tensorflow有点生疏，懒的去用啦！结果就…，线上部署必须得tensorflow模型，所以就把torch的.pt文件，转成tensorflow的pb文件把，中间遇到一些问题记录下来，仅供借鉴参考！<ps!立下flag，坚持写一些博客>
总体上pytorch转tensorflow遵循2个步骤：
这里直接利用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_names=['outputs'])
需要注意的是：tensorflow1.x模型是需要先定义计算图的，计算图需要输入输出节点，所以最好指定计算图的输入输出节点名称，以便后续调用tensorflow模型
这一步利用开源工具onnx-tf ( github) 将.onnx文件转换成.pb文件。可以直接源码安装，也可以直接 pip安装，转换过程如下：
登录后复制
def onnx2Tensorflow(onnx_model="model.onnx", tf_model="model.tf.pb"):
"""
onnx转tensorflow
"""
"""
# ---------------------报错-------------------------
# Load ONNX model and convert to TensorFlow format
model_onnx = onnx.load(onnx_model)
tf_rep = prepare(model_onnx)
# Export model as .pb file
tf_rep.export_graph(tf_model)
"""
os.system("onnx-tf convert -i %s -o %s"%(onnx_model, tf_model))
官方教程是按照注释中的内容做的，但是我自己执行报错，但是命令行方式转换可以成功。
转换成功后，我们加载计算图测试一下样例输出是否正确，如下：
登录后复制
def load_pb(path_to_pb):
"""
加载pb文件
"""
with tf.gfile.GFile(path_to_pb, 'rb') as f:
# 定义计算图
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
with tf.Graph().as_default() as graph:
tf.import_graph_def(graph_def, name='')
return graph
def loadTFModel():
"""
加载tensorflow模型测试
"""
graph = load_pb('model.tf.pb')
with tf.Session(graph=graph) as sess:
# Show tensor names in graph
for op in graph.get_operations():
# 获取节点名称
print(op.values())
# 获取输入输出节点
output_tensor = graph.get_tensor_by_name('div_3:0')
input_tensor = graph.get_tensor_by_name('inputs:0')
# dummy_input = np.random.randint(0, 1000, (1, 20), dtype=np.int64)
query1 = np.array([[4158, 7811, 6653,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]], dtype=np.int64)
query2 = np.array([[9914, 10859, 6907, 8719, 7098,
8861, 4158, 10785, 6299, 1264,
1612, 10285, 6973, 7811, 0,
0, 0, 0, 0, 0]],
dtype=np.int64)
output1 = sess.run(output_tensor,
feed_dict={input_tensor: query1})
output2 = sess.run(output_tensor,
feed_dict={input_tensor: query2})
simi = np.dot(output1[0], output2[0])
print(simi)
很明显在测试过程中，需要给输入张量喂入输出，获取输出节点数据，这个时候我们在第一步指定的输入、输出节点变量名就派上了用场。input_tensor获取输入节点的张量值就使用了变量"inputs:0"，":0"表示该节点的第一个输出，如果该节点只有一个输出就是:0;可以看到，在第一步中我们设定的输出变量名是outputs，但是这里并不是。（个人觉得是有点bug的，所以具体的输出节点名称要从计算图上的节点中去查看。当然如果使用最新版本的onnx-tf这个问题是不存在的，由于我需要的pb文件是tf 1.12版本的，亲测1.13版本以上，没有变量名映射的bug）
免责声明：本文系网络转载或改编，未找到原创作者，版权归原作者所有。如涉及版权，请联系删
155-2731-8020
