当前位置：首页 > news >正文

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

news 来源：原创 2024/4/30 15:58:08

前言

今天是程序员节，当然是以程序员的方式来度过节日。
很早就听说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