Jetson nano上部署自己的Yolov5模型（TensorRT加速）onnx模型转engine文件

背景

在主机上训练自己的Yolov5模型，转为TensorRT模型并部署到Jetson Nano上，用DeepStream运行。

硬件环境：

RTX 2080TI主机

Jetson Nano 4G B01

软件环境：

Jetson Nano:

Ubuntu 18.04

Jetpack 4.5.1

DeepStream 5.1

主机：

Ubuntu 18.04

CUDA 10.2

yolov5 5.0

训练模型（主机上）

yolov5项目链接https://github.com/ultralytics/yolov5

克隆yolov5官方的代码

git clone https://github.com/ultralytics/yolov5.git

官方训练教程详见https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data

预备环境

准备python3.8以上环境，可用conda创建一个虚拟环境，安装yolov5项目下yolov5/requirements.txt里的依赖

pip install -r requirements.txt

pytorch建议按pytorch官网教程的方式安装PyTorch

如CUDA 10.2 conda环境下安装

conda install pytorch torchvision torchaudio cudatoolkit=10.2 -c pytorch

conda安装太慢可用mamba

准备数据集

手上有一个用labelImg（https://github.com/tzutalin/labelImg）打好标签的voc格式数据集，数据集里只有两种类型，鸭子duck和马桶抽sucker。（可自行用其他方式打voc格式数据集或用labelImg直接打yolo格式数据集）

用下面代码将voc格式转成yolo格式数据集，生成images文件夹（存放所有图片），labels文件夹（存放打好的标签），test.txt（测试集），train.txt（训练集），val.txt（验证集）

import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
sets = ['train', 'test','val']
classes = ["duck", "sucker"]def convert(size, box):dw = 1. / size[0]dh = 1. / size[1]x = (box[0] + box[1]) / 2.0y = (box[2] + box[3]) / 2.0w = box[1] - box[0]h = box[3] - box[2]x = x * dww = w * dwy = y * dhh = h * dhreturn (x, y, w, h)
def convert_annotation(image_id):in_file = open('VOC2007/Annotations/%s.xml' % (image_id))out_file = open('data/labels/%s.txt' % (image_id), 'w')tree = ET.parse(in_file)root = tree.getroot()size = root.find('size')w = int(size.find('width').text)h = int(size.find('height').text)for obj in root.iter('object'):difficult = obj.find('difficult').textcls = obj.find('name').textif cls not in classes or int(difficult) == 1:continuecls_id = classes.index(cls)xmlbox = obj.find('bndbox')b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),float(xmlbox.find('ymax').text))bb = convert((w, h), b)out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
wd = getcwd()
print(wd)
for image_set in sets:if not os.path.exists('data/labels/'):os.makedirs('data/labels/')image_ids = open('VOC2007/ImageSets/Main/%s.txt' % (image_set)).read().strip().split()list_file = open('data/%s.txt' % (image_set), 'w')for image_id in image_ids:list_file.write('data/images/%s.jpg\n' % (image_id))convert_annotation(image_id)list_file.close()

按实际情况修改类型classes和文件路径。

创建配置文件

创建数据集配置文件dataset.yaml

train: data/train.txt  # 数据集里的训练集列表文件路径
val: data/val.txt  # 数据集里的验证集列表文件路径nc: 2 #类型数量names: [ 'duck', 'sucker' ]  #类型名

创建模型配置文件model.yaml，在yolov5项目下yolov5/models里复制一个要训练的模型出来修改，如yolov5s，只需修改类型数量nc

# parameters
nc: 2  # number of classes
depth_multiple: 0.33  # model depth multiple
width_multiple: 0.50  # layer channel multiple# anchors
anchors:- [10,13, 16,30, 33,23]  # P3/8- [30,61, 62,45, 59,119]  # P4/16- [116,90, 156,198, 373,326]  # P5/32# YOLOv5 backbone
backbone:# [from, number, module, args][[-1, 1, Focus, [64, 3]],  # 0-P1/2[-1, 1, Conv, [128, 3, 2]],  # 1-P2/4[-1, 3, C3, [128]],[-1, 1, Conv, [256, 3, 2]],  # 3-P3/8[-1, 9, C3, [256]],[-1, 1, Conv, [512, 3, 2]],  # 5-P4/16[-1, 9, C3, [512]],[-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32[-1, 1, SPP, [1024, [5, 9, 13]]],[-1, 3, C3, [1024, False]],  # 9]# YOLOv5 head
head:[[-1, 1, Conv, [512, 1, 1]],[-1, 1, nn.Upsample, [None, 2, 'nearest']],[[-1, 6], 1, Concat, [1]],  # cat backbone P4[-1, 3, C3, [512, False]],  # 13[-1, 1, Conv, [256, 1, 1]],[-1, 1, nn.Upsample, [None, 2, 'nearest']],[[-1, 4], 1, Concat, [1]],  # cat backbone P3[-1, 3, C3, [256, False]],  # 17 (P3/8-small)[-1, 1, Conv, [256, 3, 2]],[[-1, 14], 1, Concat, [1]],  # cat head P4[-1, 3, C3, [512, False]],  # 20 (P4/16-medium)[-1, 1, Conv, [512, 3, 2]],[[-1, 10], 1, Concat, [1]],  # cat head P5[-1, 3, C3, [1024, False]],  # 23 (P5/32-large)[[17, 20, 23], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)]

训练

从https://github.com/ultralytics/yolov5/releases下载预训练权重，如yolov5s.pt

执行yolov5项目下的train.py

python train.py --data datast.yaml --cfg model.yaml --weights yolov5s.pt --device 0

参数路径按实际情况修改

--data 数据集配置文件路径

--cfg 模型配置文件路径

--weights 预训练权重文件路径

--device CUDA设备或CPU，单显卡一般为0

其他参数详见train.py代码

运行完会在yolov5/runs/train/exp{n}/weights/下生成权重文件，best.pt为效果最佳权重，last.pt为最后一次epoch权重

用yolov5/detect.py验证识别效果

python detect.py --source path/images --weights runs/train/exp4/weights/best.pt --view-img

--source 待识别的图片文件夹路径或摄像头

--weights 权重路径

--view-img 显示识别结果

其他参数详见detect.py代码

转TensorRT

使用tensorrtx项目进行转换GitHub - wang-xinyu/tensorrtx: Implementation of popular deep learning networks with TensorRT network definition API

克隆tensorrtx项目

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

官方教程详见https://github.com/wang-xinyu/tensorrtx/tree/master/yolov5

将tensorrtx/yolov5/gen_wts.py复制到yolov5项目根目录下

执行命令生成.wts文件

python gen_wts.py yolov5/runs/train/exp4/weights/best.pt

执行完在.pt权重路径下会生成一个.wts文件

环境转到Jetson Nano

在nano上也克隆一个tensorrtx项目

将生成的.wts放到tensorrtx/yolov5/下

修改tensorrtx/yolov5/yololayer.h

static constexpr int CLASS_NUM = 2;

修改类型数量，原来是80

在tensorrtx/yolov5/目录下

编译代码

mkdir build
cd build
cmake ..
make

将.wts文件转为.engine文件

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

将验证图片放在tensorrtx/yolov5/samples/下，执行命令验证转换是否成功

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

执行完会在当前目录生成一张相应的图片

用DeepStream部署（Nano上）

安装DeepStream，DeepStream Getting Started | NVIDIA Developer

安装完后在/opt/nvidia/deepstream/deepstream-5.1/sources/objectDetector_Yolo会有一个部署yolo的官方实例代码，但只有yolov3的。

这里我们直接用已经改好的yolov5项目GitHub - DanaHan/Yolov5-in-Deepstream-5.0: Describe how to use yolov5 in Deepstream 5.0

可不用按照项目官方的说明，直接按以下步骤即可。

克隆项目

git clone https://github.com/DanaHan/Yolov5-in-Deepstream-5.0.git

进入Yolov5-in-Deepstream-5.0/Deepstream 5.0/nvdsinfer_custom_impl_Yolo/目录下

修改nvdsparsebbox_Yolo.cpp文件中的类型数量

static const int NUM_CLASSES_YOLO = 2;

原本为80，改为自己模型的类型数量

保存编译

make

返回上级目录，进入Yolov5-in-Deepstream-5.0/Deepstream 5.0/

将tensorrtx生成的.engine文件和libmyplugins.so放到目录下

这里是tensorrtx/yolov5/best.engine和tensorrtx/yolov5/builkd/libmyplugins.so

修改DeepStream处理流程配置文件deepstream_app_config_yoloV5.txt 

...[source0]
#Type - 1=CameraV4L2（usb摄像头） 2=URI（文件） 3=MultiURI
type=1
camera-width=2560
camera-height=720
camera-fps-n=30
camera-fps-d=1...[streammux]
...
width=1344
height=376
...[primary-gie]
...
model-engine-file=best.engine
...[tracker]
enable=0
tracker-width=512
tracker-height=320
ll-lib-file=/opt/nvidia/deepstream/deepstream-5.1/lib/libnvds_mot_klt.so...

我修改了[source0]，改为用摄像头实时跑识别，需要加上长宽和帧率设置。这里我用的是一个双目摄像头。

[streammux]显示窗口改为1344*376

[primary-gie]修改权重路径

[tracker]改为deepstream5.1的地址

修改config_infer_primary_yoloV5.txt文件

...model-engine-file=best.engine...num-detected-classes=2...custom-lib-path=nvdsinfer_custom_impl_Yolo/libnvdsinfer_custom_impl_Yolo.so...

修改权重路径

修改识别类型数量，原本为80

修改编译文件路径

注意，文件里还有个labels文件路径的配置

labelfile-path=labels.txt

文件夹里是没有labels.txt的，补上labels.txt标签类型文件

duck
sucker

执行命令运行DeepStream

LD_PRELOAD=./libmyplugins.so deepstream-app -c deepstream_app_config_yoloV5.txt

运行帧率在13-14fps左右

任何程序错误，以及技术疑问或需要解答的，请添加