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_demos

2、环境搭建

①参考我为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
或者
make

4、生成文件

  ① 修改脚本     

       /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

Jetson AGX Xavier 实现TensorRT加速YOLOv4_第1张图片

python trt_yolo.py --usb 0 -m yolov4-608

Jetson AGX Xavier 实现TensorRT加速YOLOv4_第2张图片

根据测试结果可知在640*480的画面中,yolov4-tiny-288模型的检测速度极快,但是检测精度较差,而且出现冗余框。yolov4-608的检测速度较原工程而言提升了四倍。

四、总结

      目标检测算法在实际应用中需要较高的实时性,但是目前很多移动嵌入式设备的算力较低,处理速度较慢,TensoRT的使用对于深度学习应用于实际而言具有很大的现实意义。

       题外话,在Xavier中跑通了YOLOv5与YOLOv4的检测代码后,个人感觉YOLOv5的整体检测效果会略优于YOLOv4。

你可能感兴趣的:(Jetson,AGX,Xavier系列,深度学习,人工智能)