标签：onnx 模型 pb Pytorch graph tensorflow tf Tensorflow model
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文件即可，基本上没有什么坑，如下：
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文件转tensorflow .pb文件
这一步利用开源工具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）
注意版本差异和环境
- 不同的tensorflow版本和onnx-tf版本需要适配：个人总结tf 1.13以上的版本可以使用最新的onnx-tf 1.7版本；tf 1.12版本可以使用onnx-tf 1.3版本。
- 注意模型推理是在cpu还是gpu上：tensorflow模型在cpu上推理和gpu上推理存在差异！torch文件默认的输入数据格式是NxCxHxW，而TensorFlow中的函数在gpu环境上都是支持NxCxHxW格式的。但是部分函数，例如max_pool在cpu 上只支持NxHxHxC的格式。解决办法就是，如果要在cpu上推理，需要禁用gpu: os.environ[“CUDA_VISIBLE_DEVICES”]="-1" ，否则转换会默认按照gpu环境导出；当然你如果没有gpu环境，那就不存在这个问题，自动按照cpu运行环境导出。
各位看客小伙伴，有什么问题或者好的解决办法可以留言评论哦！
标签：onnx,模型,pb,Pytorch,graph,tensorflow,tf,Tensorflow,model 来源： https://blog.csdn.net/qq_30666517/article/details/113249029
本站声明： 1. iCode9 技术分享网（下文简称本站）提供的所有内容，仅供技术学习、探讨和分享； 2. 关于本站的所有留言、评论、转载及引用，纯属内容发起人的个人观点，与本站观点和立场无关； 3. 关于本站的所有言论和文字，纯属内容发起人的个人观点，与本站观点和立场无关； 4. 本站文章均是网友提供，不完全保证技术分享内容的完整性、准确性、时效性、风险性和版权归属；如您发现该文章侵犯了您的权益，可联系我们第一时间进行删除； 5. 本站为非盈利性的个人网站，所有内容不会用来进行牟利，也不会利用任何形式的广告来间接获益，纯粹是为了广大技术爱好者提供技术内容和技术思想的分享性交流网站。