开服网站源码,公司网站建设的策划方案,莱芜58同城,江门网络干部学院Jetson nano上部署自己的Yolov5模型#xff08;TensorRT加速#xff09;onnx模型转engine文件 背景
在主机上训练自己的Yolov5模型#xff0c;转为TensorRT模型并部署到Jetson Nano上#xff0c;用DeepStream运行。
硬件环境#xff1a;
RTX 2080TI主机
Jetson Nano 4…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 cudatoolkit10.2 -c pytorch
conda安装太慢可用mamba
准备数据集
手上有一个用labelImghttps://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上
安装DeepStreamDeepStream 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 - 1CameraV4L2usb摄像头 2URI文件 3MultiURI
type1
camera-width2560
camera-height720
camera-fps-n30
camera-fps-d1...[streammux]
...
width1344
height376
...[primary-gie]
...
model-engine-filebest.engine
...[tracker]
enable0
tracker-width512
tracker-height320
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-filebest.engine...num-detected-classes2...custom-lib-pathnvdsinfer_custom_impl_Yolo/libnvdsinfer_custom_impl_Yolo.so...
修改权重路径
修改识别类型数量原本为80
修改编译文件路径
注意文件里还有个labels文件路径的配置
labelfile-pathlabels.txt
文件夹里是没有labels.txt的补上labels.txt标签类型文件
duck
sucker
执行命令运行DeepStream
LD_PRELOAD./libmyplugins.so deepstream-app -c deepstream_app_config_yoloV5.txt 运行帧率在13-14fps左右
任何程序错误以及技术疑问或需要解答的请添加
