本文共 1700 字，大约阅读时间需要 5 分钟。

tensorflow转trt：

# 替换不可行的算子

# 在pb代码目录，运行 CUDA_VISIBLE_DEVICES=-1 python eval_png.py 得到.pb

（在CUDA11环境下，tensorflow要在cpu上验证结果）

# 查看.pb最后一个节点

# 下载tf2onnx包：https://github.com/onnx/tensorflow-onnx

# 替换目录下tf2onnx/onnx_opset/tensor.py的内容

 

# 如果要引入自定义算子，添加目录下tf2onnx/convert.py的内容

 

# 在目录下运行 python setup.py install

# 回到pb代码目录下，修改最后一个节点（如下是Reshape_1），运行如下，生成.onnx  （有错误信息没关系，只是注册onnx算子）

python -m tf2onnx.convert --input ./FRC_1920x1080_bs1.pb --inputs "f1:0,f2:0,f1_o:0,f2_o:0,m_f1:0,m_f2:0,s_f1:0,s_f2:0" --outputs "Reshape_1:0" --output FRC_1920x1080_bs1.onnx --opset 10

# 根据CUDA版本, Ubuntu版本, 内核下载对应的TensorRT 7.2.1开源库

# 添加环境变量 export LD_LIBRARY_PATH=/home/c00443756/0result/TensorRT-7.2.1.6/lib:$LD_LIBRARY_PATH

# 在trt-C代码目录，如果有自定义算子，需要写.cpp/.cu文件，主要是写接口函数

或在trt-python代码目录，如果有自定义算子，需要写.cpp/.cu/.h文件，主要是写接口函数

# 从pb代码目录拷贝.onnx到trt-C代码目录，修改Makefile里的tensorRT路径，mkdir bin && make

问题：

Makefile.config:11: CUDA_INSTALL_DIR variable is not specified, using /usr/local/cuda by default, use CUDA_INSTALL_DIR=<cuda_directory> to change 解决： sudo make clean sudo make CUDA_INSTALL_DIR=/usr/local/cuda https://blog.csdn.net/zhw864680355/article/details/80976903

或从pb代码目录拷贝.onnx到trt-python代码目录，修改CMakeLists.txt里的tensorRT路径和.cpp/.cu/.h文件路径，mkdir build && cd build && cmake .. && make ，生成.so文件

# 在trt-C代码目录，cd bin , 运行生成的可执行文件，生成trt模型（记得先删除之前的trt模型）（该步骤在V100显卡需要半小时左右）

或在trt-python代码目录，cd ..，运行CUDA_VISIBLE_DEVICES=1 python test.py，生成trt模型（记得先删除之前的trt模型）（该步骤在V100显卡需要半小时左右）

===========================================================================================

pytorch转trt： # 固定输入尺寸，如（1,8,1080,1920）

# 直接调用torch中的onnx.export生成.onnx，有些算子可能需要opset10支持（pytorch1.1及以下没有opset_version参数）

 

# .onnx转.trt （onnx转trt叫做序列化，解析trt叫做反序列化）