基于YOLOv5框架如何训练一个自定义对象检测模型

 

 

安装与测试

最近YOLOv5最新更新升级到v6.x版本，工程简便性有提升了一大步，本教程教你基于YOLOv5框架如何训练一个自定义对象检测模型，首先需要下载对应版本：

•  

https://github.com/ultralytics/yolov5/releases/tag/v6.1

鼠标滚到最下面下载源码zip包：

•  

https://github.com/ultralytics/yolov5/archive/refs/tags/v6.1.zip

下载完成之后解压缩到：

•  

D:pythonyolov5-6.1

首先运行下面命令行完成依赖库的安装：

等待一段时间，完成全部安装之后，运行测试：

•  

python detect.py --weights yolov5s.pt --source dataimageszidane.jpg

运行结果如下：

 

数据集准备与制作

自己百度收集了一个无人机与飞鸟的数据集，其中训练集270张图像，测试集26张图像。

使用labelImg工具完成标注，工具下载地址：

•  

https://gitee.com/opencv_ai/opencv_tutorial_data/tree/master/tools

YOLOv5要求的数据集目录结构如下：

其中images是所有图像的集合，labels是所有标签信息集合。train表示训练集、valid表示测试集。这里需要注意的是，labels信息中的标签信息YOLO格式标注框需要把原始标注信息的标注框格式从：

•  

Left top right bottom

转换为

•  

Center_x, center_y, width, height

并归一化到0~1之间，这部分我写了一个脚本来完成label标签的生成，把xml的标注信息转换为YOLOv5的labels文件，这样就完成了数据集制作。最后需要创建一个dataset.ymal文件，放在与data文件夹同一层，它的内容如下：

# train and val datasets (image directory or *.txt file with image paths)train: uav_bird_training/data/images/train/val: uav_bird_training/data/images/valid/
# number of classesnc: 2
# class namesnames: ['bird', 'drone']

三：模型训练

制作好数据集之后，模型训练就成为一件很简单事情，一条命令行搞定。运行下面的命令行：

•  

python train.py --img 640 --batch 4 --epochs 25 --data uav_bird_trainingdataset.yaml --weights yolov5s.pt

其中uav_bird_training文件夹里是制作好的数据集。这样就开始训练，训练过程中可以通过tensorboard来查看可视化的结果，

 

PR曲线说明训练效果还错！

不同框架与硬件平台推理比较

YOLOv5的6.x版本支持不同框架模型导出与推理，看下图：

这里分别导出模型为ONNX、XML、engien格式之后，使用同一段视频在我的笔记本上（CPUi7, GPU3050ti）：在下面四种推理框架上运行

-OpenCV DNN-OpenVINO-ONNXRUNTIME-TensorRT

测试，分别截图如下：

OpenCV DNN推理速度

 

OpenVINO平台上的推理速度

 

ONNXRUNTIME GPU推理速度

 

TensorRT框架部署-FP32版本模型推理统计：

 

YOLOv5的6.x版本，是支持TensorRT 量化到FP16模型直接导出的，但是不支持INT8量化生成，所以自己实现了导出 量化INT8版本，测试结果如下：

•  

python detect.py --weights uav_bird_training/uav_bird_int8.engine --data uav_bird_training/dataset.yaml --source D:/bird/bird_fly.mp4

 

TensorRT框架部署-INT8版本模型推理统计：

 

最终比较：

注意：以上结果是基于Python语言上测试结果（前后处理都使用pytorch完成）  

 

番外篇：C++推理与比较

使用C++部署时候，前后处理都改成了基于OpenCV 完成，使用CPU完成前后处理，OpenVINO+CPU运行速度截图如下：

 

OpenCV DNN + CUDA版本推理

 

TensorRT-FP32模型推理速度

 

TensorRT-INT8模型推理速度

 

最后把自定义训练导出的模型转换为TensorRT Engine文件之后，部署到了我的一块Jetson Nano卡上面，实现了边缘端的部署，有图有真相：

 

  审核编辑 ：李倩