TensorRt推理部署优化方案及流程概述

1.参考资料

先放上参考资料，以表感谢。

TensorRT-Installation-Guide.pdf
TensorRT-Developer-Guide.pdf (PG-08540-001_v8.4.1 | June 2022)
TensorRT-Sample-Support-Guide.pdf
CSDN:扫地的小何尚
github:TensorRT
https://pytorch.org/tutorials/

2.流程概述

1.首先，进行模型转换。目前掌握的模型转换方案有两种，分别是pytorch–>onnx–>engine方案和pytorch–>wts–>engine方案。

2.使用转换得到的TensorRT模型进行推理

3.TensorRT推理结果的decode和后处理

3.pytorch–>onnx–>engine方案

3.1 pytorch–>onnx转换

onnx是一种比较通用的模型表示形式，它不仅可以保存模型的网络结构，还可以保存模型的参数，因此将模型从pytorch导出到onnx后，不仅不需要pytorch权重文件（即：pth文件）了，而且使用pytorch定义的网络结构也不需要了。因为在导出时这些信息都已经保存到了onnx文件中去了。

这一步使用pytorch框架中的接口完成，一般是在将图像输入网络之前调用torch.onnx.export接口，因为这样代码改动最少，更容易的找到网络的输入张量的维度和加载好的pytorch模型。

代码实现细节参考Pytorch官方教程链接。

3.2 onnx–>engine转换

在进行如下操作前需要安装好TensorRT，可以参考官方教程:TensorRT-Installation-Guide.pdf.

不同于pytorch–>onnx，而onnx–>engine转换具有多种实现方式。大概有三种方法：
1.使用TensorRT python API
2.使用TensorRT C++ API
3.使用TensorRT trtexec 可执行文件

python API和C++ API的流程是一样的，但是C++ API需要编写CMakeLists.txt或者Makefile文件，代码复杂度也更高一些，因此更建议使用python API的方式。python API的代码实现请参考TensorRT-Developer-Guide.pdf 4.1.2节。

trtexec适用于使用.tar.gz压缩文件的方法安装TensorRT时，解压后在Samples中有一个trtexec的项目，编译之后便可以在bin目录下生成一个trtexec的可执行文件。trtexec的使用可以参考官方文档TensorRT-Developer-Guide.pdf的附录部分，或者CSDN该中文资料链接。

3.3 engine推理

使用转换的engine模型进行推理，这里特别说明本文推理指的是网络的forward过程。

首先从文件中读取engine文件，进行反序列化，生成engine模型，然后在engine在基础上构造一个context，后面的推理使用的便是此context。

使用engine进行推理的具体细节可以参考：TensorRT-Developer-Guide.pdf. C++ API 3.2和3.4节, Python API 4.2和4.4节。

本文所讲的推理，输出是网络forward的原始结果，以我的一个车道线检测的算法为例，输出的两个数组如下图所示：

该网络输出两个数组是因为该网络有两个输出头，分别对应是车道线关键点和车道线属性。

3.4 推理结果decode

如上节所示，TensorRT推理出来的结果无法直接使用，还需要一个decode的过程，一般是按照算法的原理或仿照Pytorch框架的推理逻辑，将TensorRT推理得到的数组decode为合适的张量或多维数组。以上面的算法为例，第一个25728维的数组需要解码为(1, 201, 32, 4)的一个张量或多维数组，第二个32位的数组需要解码为(1, 2, 16)的张量或多维数组。

至此，TensorRT的推理结果在形式上与Pytorch的推理结果可以实现对应了，在这一步可以对比两者的结果大致判断TensorRT的推理是否准确。然后，还需要后处理的一些操作，这部分与实际的功能需求而定，无论是TensorRT还是Pytorch，它们后处理的算法逻辑基本一致，只不过是代码实现上的差别。相比之下TensorRT的后处理要更加注重代码的时间复杂度。

4.pytorch–>wts–>engine方案

4.1 yolov5:pytorch–>wts

这一步转换只是将Pytorch的模型和参数读取并按照固定的格式写入到wts文件中，参考tensorrtx的gen_wts.py脚本即可。

4.2 wts–>engine

1.使用TensorRT重新实现与Pytorch模型一样的网络结构，使用TensorRT定义网络并导入权重的代码实现可以参考TensorRT-Developer-Guide.pdf 6.4节。

2.构建后网络后，需要生成engine，这一步转换可以是在线的也可以是离线完成的。如果是在线的，则每次程序运行时，都需要进行一次转换。如果是离线的可以将生成的engine保存下来，程序以后每次运行时仅需从engine文件反序列化即可。该方案的其余工作与pytorch–>onnx–>engine方案基本相同，不再重复。