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.txtpytorch建议按pytorch官网教程的方式安装PyTorch
如CUDA 10.2 conda环境下安装
conda install pytorch torchvision torchaudio cudatoolkit=10.2 -c pytorchconda安装太慢可用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左右
任何程序错误，以及技术疑问或需要解答的，请添加
