【Jetson平台 ubuntu 18.04 使用 TensorRT 加速 YOLOv5 检测】

Pytorch版本的模型推理速度有限，为了提高检测速率，使用TensorRT部署。
注意：tensorrt转换出来的模型实际上是和硬件绑定的，也就是在部署的过程中，如果你的显卡和显卡相关驱动软件（cuda、cudnn）发生了改变，那么模型就得需要重新做转换。

1）克隆TensorRT代码

首先克隆TensorRT编译YOLOv5的代码

git clone https://github.com/wang-xinyu/tensorrtx.git  

可以看到TensorRT文件夹中包含编译多种神经网络的文件，只保留yolov5文件夹，如图：

图22- 38 TensorRT文件夹

2）模型转换

在tensorrtx/yolov5文件夹中可找到gen_wts.py，该脚本可将Pytorch模型(.pt格式)转换成权重文本文件(.wts格式)。将gen_wts.py拷贝到{ultralytics}/yolov5文件夹中，并在此文件夹中打开终端，输入

python gen_wts.py -w best.pt -o 713best.wts

``
图22- 39 模型转换
这样就在{ultralytics}/yolov5文件夹中生成了best.wts文件。

3）修改yololayer.h参数

打开tensorrtx\yolov5文件夹，找到yololater.h文件，修改CLASS_NUM，使其与pytorch模型一致：
图22- 40 修改yololayer.h参数

4）编译并生成引擎文件

终端依次输入下列指令，可在tensortrtx/yolov5目录下新建build文件夹

cd work/tensorrtx/yolov5   
mkdir bulid  

将上一步生成的713best.wts拷贝到build文件夹下，再依次输入下列指令对其进行编译

cmake ..  
make  

图22- 41 编译713best.wts文件_cmake

图22- 42 编译713best.wts文件_make
再输入下列指令，可使用713best.wts生成713best.engine，注意末尾的n代表模型尺度，可选项有n/s/m/l/x/n6/s6/m6/l6/x6 or c/c6 gd gw，必须与pytorch模型一致：

sudo ./yolov5 -s 713best.wts 713best.engine n 

图22- 43 生成engine文件

5）进行检测

输入下列指令，可对指定路径下的图片进行检测，这里的图片路径是…/samples：

sudo ./yolov5 -d 713best.engine ../samples  

图22- 44 图片检测
与使用Pytorch模型检测所耗时间（55ms_{80ms）进行对比，TensorRT加速后单幅图片耗时明显缩短（17ms}27ms）。

图22- 45 检测结果