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

桃扇骨 2022-09-05 00:17 697阅读 0赞

## 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][https_github.com_ultralytics_yolov5]

克隆yolov5官方的代码

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

官方训练教程详见[https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data][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/][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][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][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][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][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

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

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2FpbGFpZXI_size_16_color_FFFFFF_t_70][]

# 用DeepStream部署（Nano上） #

安装DeepStream，[https://developer.nvidia.com/deepstream-getting-started][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][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

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2FpbGFpZXI_size_16_color_FFFFFF_t_70 1][]

运行帧率在13-14fps左右

[https_github.com_ultralytics_yolov5]: https://github.com/ultralytics/yolov5
[https_github.com_ultralytics_yolov5_wiki_Train-Custom-Data]: https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data
[https_pytorch.org]: https://pytorch.org/
[https_github.com_tzutalin_labelImg]: https://github.com/tzutalin/labelImg
[https_github.com_ultralytics_yolov5_releases]: https://github.com/ultralytics/yolov5/releases
[https_github.com_wang-xinyu_tensorrtx]: https://github.com/wang-xinyu/tensorrtx
[https_github.com_wang-xinyu_tensorrtx_tree_master_yolov5]: https://github.com/wang-xinyu/tensorrtx/tree/master/yolov5
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2FpbGFpZXI_size_16_color_FFFFFF_t_70]: /images/20220829/6b0fd224070e40f0ac31a881d5ae4da7.png
[https_developer.nvidia.com_deepstream-getting-started]: https://developer.nvidia.com/deepstream-getting-started
[https_github.com_DanaHan_Yolov5-in-Deepstream-5.0]: https://github.com/DanaHan/Yolov5-in-Deepstream-5.0
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2FpbGFpZXI_size_16_color_FFFFFF_t_70 1]: /images/20220829/ea0d78a29cc04cfeaed3c112f0b8b114.png