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

1.参考资料

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

  • TensorRT-Installation-Guide.pdf
  • TensorRT-Developer-Guide.pdf (PG-08540-001_v8.4.1 | June 2022)
  • TensorRT-Sample-Support-Guide.pdf

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的附录部分。

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方案基本相同,不再重复。