Jetson AGX Xavier 实现TensorRT加速YOLOv4
上一篇：Jetson AGX Xavier测试YOLOv4
下一篇：Jetson AGX Xavier安装ROS
一、前言
前文，在Xavier中测试YOLOv4算法发现其检测速度较慢，可以采用TensorRT对其进行加速。目前很多大佬都对其进行了实现。本文采用Github 中 JK Jung 作者的tensorrt_demos工程来提高yolov4的检测速度。
二、加速准备
1、克隆tensorrt_demos工程
git clone https://github.com/jkjung-avt/tensorrt_demos2、环境搭建
①参考我为yolov5配置环境及TensorRT加速的文章，配置的python环境，得到名为yolov5env的环境。
②安装pycuda与onnx
conda activate yolov5env
pip install pycuda
pip install onnx==1.4.1切记不要用最新版的onnx，会报错。这里pycuda的安装可能会出现问题，建议采用源码安装。我也忘记当初是如何解决，只记得安装失败过。网上挺多相关教程。
3、编译工程
进入/tensorrt_demos/plugins/ 目录下：
make -j8
或者
make4、生成文件
① 修改脚本
/tensorrt_demos/yolo/目录下我们会发现download_yolo.sh脚本，脚本打开我们会发现该脚本会下载包括YOLOv3与YOLOv4中许多版本的cfg文件与权重文件，并生成不同size的相关文件。但是有些文件我并不需要，而且国内网络容易被墙或者下载较慢。
为了方便测试，本文主要以yolov4-tiny.weights与yolov4.weights为例。故删除了download_yolo.sh中的部分指令 ，我们先将预先下载好的yolov4-tiny.weights，yolov4-tiny.cfg ， yolov4.weights，yolov4.cfg文件（文件有需要的可以私信我，不知道为啥上传CSDN的文件会被自动收费。）放入/tensorrt_demos/yolo/目录下。然后将download_yolo.sh脚本内容修改如下：
#!/bin/bash
set -e
echo
echo "Creating yolov4-tiny-288.cfg and yolov4-tiny-288.weights"
cat yolov4-tiny.cfg | sed -e '6s/batch=64/batch=1/' | sed -e '8s/width=416/width=288/' | sed -e '9s/height=416/height=288/' > yolov4-tiny-288.cfg
echo >> yolov4-tiny-288.cfg
ln -sf yolov4-tiny.weights yolov4-tiny-288.weights
echo "Creating yolov4-tiny-416.cfg and yolov4-tiny-416.weights"
cat yolov4-tiny.cfg | sed -e '6s/batch=64/batch=1/' > yolov4-tiny-416.cfg
echo >> yolov4-tiny-416.cfg
ln -sf yolov4-tiny.weights yolov4-tiny-416.weights
echo "Creating yolov4-288.cfg and yolov4-288.weights"
cat yolov4.cfg | sed -e '2s/batch=64/batch=1/' | sed -e '7s/width=608/width=288/' | sed -e '8s/height=608/height=288/' > yolov4-288.cfg
ln -sf yolov4.weights yolov4-288.weights
echo "Creating yolov4-416.cfg and yolov4-416.weights"
cat yolov4.cfg | sed -e '2s/batch=64/batch=1/' | sed -e '7s/width=608/width=416/' | sed -e '8s/height=608/height=416/' > yolov4-416.cfg
ln -sf yolov4.weights yolov4-416.weights
echo "Creating yolov4-608.cfg and yolov4-608.weights"
cat yolov4.cfg | sed -e '2s/batch=64/batch=1/' > yolov4-608.cfg
ln -sf yolov4.weights yolov4-608.weights
echo
echo "Done."② 执行脚本生成相关文件。
sh download_yolo.sh
完成后会生成以下文件：
③ 生成onnx文件与tensorrt引擎文件
在/tensorrt_demos/yolo/目录下打开终端：
conda activate yolov5env以yolov4-tiny-288为例：
python yolo_to_onnx.py -m yolov4-tiny-288执行以上代码会生成 yolov4-tiny-288.onnx文件。
接着执行：
python onnx_to_tensorrt.py -m yolov4-tiny-288生成yolov4-tiny-288.trt文件。到此整个加速准备步骤全部结束。
根据以上示例，可以生成yolov4-tiny-416.trt , yolov4-208.trt , yolov4-416.trt , yolov4-608.trt。等文件。
三、加速实现
本文以yolov4-tiny模型为例，进入/tensorrt_demos/yolo/ 目录下。打开终端。
conda activate yolov5env
python trt_yolo.py --usb 0 -m yolov4-tiny-288
python trt_yolo.py --usb 0 -m yolov4-608
根据测试结果可知在640*480的画面中，yolov4-tiny-288模型的检测速度极快，但是检测精度较差，而且出现冗余框。yolov4-608的检测速度较原工程而言提升了四倍。
四、总结
目标检测算法在实际应用中需要较高的实时性，但是目前很多移动嵌入式设备的算力较低，处理速度较慢，TensoRT的使用对于深度学习应用于实际而言具有很大的现实意义。
题外话，在Xavier中跑通了YOLOv5与YOLOv4的检测代码后，个人感觉YOLOv5的整体检测效果会略优于YOLOv4。