Jetson nano上部署自己的Yolov5模型（TensorRT加速）

背景

在主机上训练自己的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项目链接GitHub - ultralytics/yolov5: YOLOv5 🚀 in PyTorch > ONNX > CoreML > TFLite

克隆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（GitHub - heartexlabs/labelImg: LabelImg is now part of the Label Studio community. The popular image annotation tool created by Tzutalin is no longer actively being developed, but you can check out Label Studio, the open source data labeling tool for images, text, hypertext, audio, video and time-series data.）打好标签的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.0
    y = (box[2] + box[3]) / 2.0
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return (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').text
        cls = obj.find('name').text
        if cls not in classes or int(difficult) == 1:
            continue
        cls_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)
  ]

训练

从Releases · ultralytics/yolov5 · GitHub下载预训练权重，如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项目进行转换https://github.com/wang-xinyu/tensorrtx

克隆tensorrtx项目

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

官方教程详见tensorrtx/yolov5 at master · wang-xinyu/tensorrtx · GitHub

将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 SDK - Get Started | NVIDIA Developer

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

这里我们直接用已经改好的yolov5项目https://github.com/DanaHan/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左右