最近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
下载完成之后解压缩到: python detect.py --weights yolov5s.pt --source data\images\zidane.jpg 自己百度收集了一个无人机与飞鸟的数据集,其中训练集270张图像,测试集26张图像。使用labelImg工具完成标注,工具下载地址: https://gitee.com/opencv_ai/opencv_tutorial_data/tree/master/tools
YOLOv5要求的数据集目录结构如下: 其中images是所有图像的集合,labels是所有标签信息集合。train表示训练集、valid表示测试集。这里需要注意的是,labels信息中的标签信息YOLO格式标注框需要把原始标注信息的标注框格式从:转换为 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 classes nc: 2
# class names names: ['bird', 'drone'] 制作好数据集之后,模型训练就成为一件很简单事情,一条命令行搞定。运行下面的命令行:python train.py --img 640 --batch 4 --epochs 25 --data uav_bird_training\dataset.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完成)想要获得YOLOv5从数据集制作到训练到四种平台OpenCV DNN, OpenVINO, ONNXRUNTIME, TensorRT部署的全部技能,解锁深度学习各种平台部署技巧与代码实现,本教程全部源码与相关PPT,立刻扫码:使用C++部署时候,前后处理都改成了基于OpenCV 完成,使用CPU完成前后处理,OpenVINO+CPU运行速度截图如下:
OpenCV DNN + CUDA版本推理 TensorRT-FP32模型推理速度 TensorRT-INT8模型推理速度 最后把自定义训练导出的模型转换为TensorRT Engine文件之后,部署到了我的一块Jetson Nano卡上面,实现了边缘端的部署,有图有真相:
|