2022-4-14如何将基于tensorflow的.pb模型转成.trt模型

如何将基于tensorflow的.pb模型转成.trt模型

1.将.pb模型转成.onnx模型

步骤一:查看.pb模型的节点名称
方法①:直接使用现成的工具进行查看,推荐netron,链接网址:https://netron.app/
该工具可以直观的看到模型的结构,以及输入输出的节点名称
参考图如下(只截取了一部分):
在这里插入图片描述
方法②:写程序查看.pb模型的节点名称

import tensorflow as tf
with tf.compat.v1.Session() as sess:
    with open('./xxxxx.pb', 'rb') as f:
        graph_def=tf.compat.v1.GraphDef()
        graph_def.ParseFromString(f.read())
        print(graph_def)

注意:tfcompat.v1.Session()一定要加(),否则会出现以下问题:
在这里插入图片描述
运行该程序可以获取.pb模型的输入输出节点名称,参考结果如下(只截取部分结果):
在这里插入图片描述
步骤二:将.pb模型转成.onnx模型
直接使用tf2onnx工具便可以完成转换

python3 -m tf2onnx.convert \
              --input ./xxxx.pb \
              --inputs images:0 \ 
              --outputs features:0 \ 
              --output ./xxxx.onnx \ 
              --verbose

注意: inputs和outputs要写成node_name:port_id的形式(我个人对port_id的理解是指输入的和输出的通道,比如说多输入多输出时,会涉及到多个port_id),如若理解有误,还望大家不吝赐教(我也是一个还在慢慢探索深度学习相关内容的学习者,所以非常希望大家说出我的问题哈哈)
运行该代码可能会出现以下问题:
在这里插入图片描述
此时需要在源代码做一点修改(具体原因在参考链接里面有详细说明,我就不再赘述):

python3 -m tf2onnx.convert \
              --input ./xxxx.pb \
              --inputs images:0 \ 
              --outputs features:0 \ 
              --output ./xxxx.onnx \ 
              --opset 10 \
              --verbose

运行代码即可得到.onnx模型

2.将.onnx模型转成.trt模型

该步骤可参见博客:https://blog.csdn.net/LJ1120142576/article/details/121904431
①若出现以下问题:
在这里插入图片描述
可以按照链接: link.给的方法进行解决,然后在将.onnx模型转成.trt模型

参考博客链接
1)https://blog.csdn.net/m0_51004308/article/details/116152611
2)https://blog.csdn.net/weixin_34307464/article/details/85971793
3)https://blog.csdn.net/qq_39678190/article/details/105481102
4)https://blog.csdn.net/sinat_37532065/article/details/106664724

猜你喜欢

转载自blog.csdn.net/LJ1120142576/article/details/123106186