TensorRT来加速YOLO v5推理与检测

TensorRT来加速YOLO v5推理与检测

TensorRT简介

TensorRT是一个高性能的深度学习推理（Inference）优化器，可以为深度学习应用提供低延迟、高吞吐率的部署推理。TensorRT可用于对超大规模数据中心、嵌入式平台或自动驾驶平台进行推理加速。TensorRT现已能支持TensorFlow、Caffe、Mxnet、Pytorch等几乎所有的深度学习框架，将TensorRT和NVIDIA的GPU结合起来，能在几乎所有的框架中进行快速和高效的部署推理。
TensorRT 是一个C++库，从 TensorRT 3 开始提供C++ API和Python API，主要用来针对 NVIDIA GPU进行 高性能推理（Inference）加速。

基本特性：

    用于高效实现已训练好的深度学习模型的推理过程的SDK

    内含推理优化器和运行时环境

    使DL模型能以更高吞吐量和更低的延迟运行

    有C++和python的API，完全等价可以混用

TensorRT主要做了以下几点来加快推理速度：

    算子融合

    量化

    动态张量显存
    
    内核自动调整

    多流执行

一、TensorRT安装

1. 电脑基础环境

• 操作系统: win10
• 显卡： 4060TI 16GB
• pytorch 版本：1.13.0 使用 print(torch._ _version_ _) 查看 pytorch 版本
• python: 3.8

2. 查看 cuda 版本

• 激活 conda 中的一个环境, 并输入 nvcc -V 进行查看, 我的版本如下图所示
  

3. 安装TensorRT

需要到nvidia官网去下载对应cuda版本的TensorRT安装包,我这里下载对应版本压缩包如下，下载解压即可。

下载好将压缩包进行解压，激活 conda 环境, 然后 cd 进入下面几个文件夹，安装 whl 文件。

• 将 解压后的TensorRT-8.5.3.1\include 中头文件拷贝到 安装的cuda目录中 C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.7\include中
• 将 解压后的TensorRT-8.5.3.1\lib 中所有的lib文件拷贝到 安装的cuda目录中 C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.7\lib\x64中
• 将 解压后的TensorRT-8.5.3.1\lib 中所有的dll文件拷贝到 安装的cuda目录中 C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.7\bin中
• 将 :\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.7\bin 路径添加到用户的 Path 环境变量中
  
  在上面操作添加以后，重启电脑，此时环境变量生效。

4. 验证TensorRT

激活conda 环境，依次输入下面语句。

python
import tensorrt
print(tensorrt.__version__)

并运行下载解压后安装包里的一个实例的 sample.py 文件, 来检查 tensorrt 是否能正常使用。
直接在激活的环境下， cd 进入，比如我的文件目录， 直接cd
C:\Users\xiaoalei\Desktop\TensorRT-8.5.3.1\samples\python\network_api_pytorch_mnist

• 主要看后面的一部分地址
  
  -输入 python sample.py 运行
  
• 出现上面情况就是表示成功，（运行报错一般是安装包没有安装，直接 pip install 就行，我在安装提示只需安装 pycuda，直接 pip install pycuda）

二、YOLO v5模型导出函数

而在YOLOv5 6.0版本更新后，新增了一个export.py文件，支持大部分框架模型的导出，包括TensorRT。这里介绍部分模型导出函数

1.onnx 模型导出（def export_onnx（）函数）

2.openvino模型导出（def export_openvino（）函数）

3.coreml模型导出（def export_coreml（）函数）

4.TensorRT模型导出（def export_engine（）函数）

三、YOLO v5使用TensorRT进行加速检测

1. 模型导出

直接导出模型：（我这里对自己训练好的模型和权重文件进行导出）

• 这里的--half表示半精度模型，使用半精度可以加快推理速度，但会损失一定精度，直接导出可以不加

 python export.py --weights runs/train/exp/weights/best.pt --data linhuo.yaml --include engine --device 0 --half

2. 图片检测

验证图片检测速度, 需要对 detect.py 进行修改，最后的程序入口进行修改
修改后如下：

if __name__ == '__main__':

    start_time = time.time()
    opt = parse_opt()
    main(opt)
    end_time = time.time()
    print("程序花费时间{}秒".format(end_time - start_time))

3. 使用yolov5s.pt 和 yolov5s.engine（tensorrt）进行检测

命令行语句：

3.1 使用yolov5s.pt, 推理时间: 5.49s

python detect.py --weights runs/train/exp/weights/best.pt

3.2 使用 yolov5s.engine（tensorrt）, 推理时间: 4.20s

python detect.py --weights runs/train/exp/weights/best.engine

四、总结

TensorRT能加速YOLO v5推理与检测。
后续会补充视频推理检测