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项目链接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官网教程的方式安装https://pytorch.org/

如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.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)
  ]

训练

从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项目进行转换https://github.com/wang-xinyu/tensorrtx

克隆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，https://developer.nvidia.com/deepstream-getting-started

安装完后在/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左右