【目标检测】使用TensorRT加速YOLOv5

前言

今天是程序员节，当然是以程序员的方式来度过节日。 很早就听说TensorRT可以加速模型推理，但一直没时间去进行实践，今天就来把这个陈年旧坑填补一下。

背景知识

在实践之前有必要了解一下相关知识。

TensorRT简介

TensorRT是可以在NVIDIA各种GPU硬件平台下运行的一个C++推理框架。我们利用Pytorch、TF或者其他框架训练好的模型，可以转化为TensorRT的格式，然后利用TensorRT推理引擎去运行我们这个模型，从而提升这个模型在英伟达GPU上运行的速度[1]。

TensorRT支持几乎所有主流深度学习框架，将python框架转换成C++的TensorRT，从而可以加速推理。

具体而言，TensorRT主要做了以下几点来加快推理速度[1]：

• 算子融合(层与张量融合)：简单来说就是通过融合一些计算op或者去掉一些多余op来减少数据流通次数以及显存的频繁使用来提速
• 量化：量化即IN8量化或者FP16以及TF32等不同于常规FP32精度的使用，这些精度可以显著提升模型执行速度并且不会保持原先模型的精度
• 内核自动调整：根据不同的显卡构架、SM数量、内核频率等(例如1080TI和2080TI)，选择不同的优化策略以及计算方式，寻找最合适当前构架的计算方式
• 动态张量显存：我们都知道，显存的开辟和释放是比较耗时的，通过调整一些策略可以减少模型中这些操作的次数，从而可以减少模型运行的时间
• 多流执行：使用CUDA中的stream技术，最大化实现并行操作

当然，TensorRT主要缺点是与特定GPU绑定[1]，在不同型号上转换出来的模型不能通用(这一点笔者暂未去从实践证实)

TensorRT官方在其仓库提供了三个开源工具，之后有需要可以使用。

三个工具大致用途[1]：

• ONNX GraphSurgeon 可以修改我们导出的ONNX模型，增加或者剪掉某些节点，修改名字或者维度等等
• Polygraphy 各种小工具的集合，例如比较ONNX和trt模型的精度，观察trt模型每层的输出等等，主要用-来debug一些模型的信息
• PyTorch-Quantization 可以在Pytorch训练或者推理的时候加入模拟量化操作，从而提升量化模型的精度和速度，并且支持量化训练后的模型导出ONNX和TRT

实践上手

下面来进行实践上手。

实验环境

这是我的电脑环境：

• 操作系统：Windows10
• 显卡：RTX2060
• CUDA版本：11.6
• Pytorch版本：1.7.1
• Python版本：3.8

查看CUDA版本

TensortRT非常依赖CUDA版本，在安装之前，需要先查看本机安装好的CUDA版本，查看方式有多种，第一种方式可以通过NVIDIA 控制面板查看；第二种方式可以通过在控制台输入nvcc -V进行查看。

安装TensorRT

首先需要到Nvidia官网去下载对应Cuda版本的TensorRT安装包。

我这里下载的是红框选中的这一个，这个版本支持CUDA11.0-11.7

下载好后，将压缩包进行解压缩，然后到Anaconda环境中，以此进入到以下几个文件夹，安装这几个whl文件

cd TensorRT-8.4.3.1\python
pip install tensorrt-8.4.3.1-cp38-none-win_amd64.whl

cd TensorRT-8.4.3.1\graphsurgeon
pip install graphsurgeon-0.4.6-py2.py3-none-any.whl

cd TensorRT-8.4.3.1\onnx_graphsurgeon
pip install onnx_graphsurgeon-0.3.12-py2.py3-none-any.whl

cd TensorRT-8.4.3.1\uff
pip install uff-0.6.9-py2.py3-none-any.whl

然后需要移动安装包里的一些文件：

• 将TensorRT-8.4.3.1\include中头文件拷贝到C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.6\include
• 将TensorRT-8.4.3.1\lib中所有lib文件拷贝到C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.6\lib\x64
• 将TensorRT-8.4.3.1\lib中所有dll文件拷贝到C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.6\bin

注：这里的v11.6根据自己的Cuda版本号即可

之后，需要手动将C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.6\bin路径添加到用户Path环境变量中

然后进行验证：

python
import tensorrt
print(tensorrt.__version__)

在import时发生报错，

FileNotFoundError: Could not find: nvinfer.dll. Is it on your PATH?

此时只需要将缺少的文件找到，然后添加到上面的bin目录下即可，我这里是在安装的torch中lib文件下找到的部分文件，缺什么移什么即可。

如无报错，再次验证，可以输出tensorrt版本：

下面运行安装包里面的一个sample.py文件，以确保tensorrt能够正常工作。 进入到下图所示的路径，运行sample.py，如果正常输出，则代表tensorrt安装成功。

如果提示没装pycuda，还需要再安装一下

pip install pycuda

YOLOv5使用TensorRT加速

下面内容有点乱，是我实验时的草稿，明天再起来整理。

python export.py --weights yolov5s.pt --data data/coco128.yaml --include engine --device 0 --half

ONNX: export failure 0.4s: Exporting the operator silu to ONNX opset version 12 is not supported.

pytorch的激活函数源码位置：X:\anaconda3\envs\python3.7\Lib\sitepackages\torch\nn\modules\activation.py(此处结合自己的anaconda实际安装位置来更改)。

class SiLU(Module):

    __constants__ = ['inplace']
    inplace: bool
 
    def __init__(self, inplace: bool = False):
        super(SiLU, self).__init__()
        self.inplace = inplace
 
    def forward(self, input: Tensor) -> Tensor:
        # ------------------------------------- #
        # 把F.silu替换掉，修改后如下
        return input * torch.sigmoid(input)
 
        #原来的代码
        return F.silu(input, inplace=self.inplace)

实验结果

导出花费时间

Export complete (404.2s) 半精度

不是半精度 Export complete (99.5s)

80多秒

实验

修改detect

if __name__ == "__main__":
    begin_time = time.time()
    opt = parse_opt()
    main(opt)
    end_time = time.time()
    print("程序花费时间{}秒".format(end_time-begin_time))

执行

python detect.py --weights yolov5s.pt

程序花费时间8.395954370498657秒

半精度

python detect.py --weights yolov5s.engine

程序花费时间4.830101728439331秒

全精度

程序花费时间5.456318378448486秒

计算一下速度提升35%

pt all 128 929 0.719 0.625 0.714 0.475

半精度 all 128 929 0.682 0.64 0.708 0.47

全精度 all 128 929 0.677 0.659 0.711 0.472

检测视频

正常检测 程序花费时间20.880696058273315秒

转换成tensorrt 程序花费时间16.03518009185791秒

转成onnx推理 程序花费时间25.689934253692627秒

References

[1]TensorRT详细入门指北：https://blog.csdn.net/IAMoldpan/article/details/117908232 [2]https://blog.csdn.net/weixin_43917589/article/details/122578198 [3]https://blog.csdn.net/qq_44523137/article/details/123178290?spm=1001.2014.3001.5502