YOLOv6算法新鲜出炉--训练自己数据集过程


前言

提示;本文记录的是使用YOLOv6算法训练自己数据集的过程:

YOLOv6算法背景:YOLOv6 是美团视觉智能部研发的一款目标检测框架,致力于工业应用。本框架同时专注于检测的精度和推理效率,在工业界常用的尺寸模型中:YOLOv6-nano 在 COCO 上精度可达 35.0% AP,在 T4 上推理速度可达 1242 FPS;YOLOv6-s 在 COCO 上精度可达 43.1% AP,在 T4 上推理速度可达 520 FPS。
评估指标如下;
YOLOv6算法新鲜出炉--训练自己数据集过程_第1张图片

官方公众哈介绍说 YOLOv6算法精度与速度远超 YOLOv5 和 YOLOX 的新框架,是真是假,还待自己跑跑测评一下,下面将介绍如何使用YOLOv6算法训练自己的数据集。
源码链接:https://github.com/meituan/YOLOv6

一、环境搭建

本人搭建的环境:Pytorch1.8.1 ,torchvision0.9.1, python3.8.1 cuda11.1 ,显卡是2080TI,系统为ubuntu18.04

代码下载后的结构如下图:
YOLOv6算法新鲜出炉--训练自己数据集过程_第2张图片
安装requirements.txt中的安装包,可以使用清华源快速安装,安装命令:

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple -r requirements.txt

如果出现个别安装不上的,建议单独通过pip install xxx来安装,requirements.txt内容如下:
YOLOv6算法新鲜出炉--训练自己数据集过程_第3张图片
博主在安装pycocotools出现找不到的情况,然后就直接使用pip install pycocotools来安装,只是安装过程较慢,需要等待。如果您是windows系统下安装pycocotools,好像麻烦点,需要依赖其他,不过都有博客介绍。

二、数据集制作

1.数据标注格式为yolo格式,跟yolov5使用的数据格式一模一样

2.data文件夹下仿照coco.yaml新建一个自己的.yaml文件

博主训练的是电单车、人两类目标的数据集,新建的yaml命名为:person_ebike.yaml
.yaml中的内容,包含训练集、验证集图片存放的文件夹的路径、类别数、类别名字。

值得注意的是文件夹路径最后不要带“/”,否则训练会报错,博主遇到坑,因为yolov5算法yaml文件中的图片路径最后是可以带“/”。
如正确的yaml,自建的yaml文件中只保留train:和val:就可以,那个test和anno_path删除。
YOLOv6算法新鲜出炉--训练自己数据集过程_第4张图片
路径结尾带有“/”,错误的yaml
YOLOv6算法新鲜出炉--训练自己数据集过程_第5张图片

配置好自己训练数据集对应的yaml文件后,就可以开始训练了,博主使用的是yolov6s.pt作为预训练模型。无需改其他文件,很简单。无非在train.py中设置epoch、batchsize、输入图片大小等。

三、模型训练

在YOLOv6文件夹下执行命令开始单卡训练:

python tools/train.py --batch 32 --conf configs/yolov6s_finetune.py --data data/person_ebike.yaml --device 0
![在这里插入图片描述](https://img-blog.csdnimg.cn/e448eeecac734aabaec25a92e6c5f1d9.jpeg#pic_center)

训练过程如下图所示:
YOLOv6算法新鲜出炉--训练自己数据集过程_第6张图片

YOLOv6与YOLOv5效果测试对比

分别用电梯内电单车+人形数据集训练两种算法的模型,模型迭代200次。
YOLOv5模型实际测试效果如下:
YOLOv6算法新鲜出炉--训练自己数据集过程_第7张图片
YOLOv6算法新鲜出炉--训练自己数据集过程_第8张图片
YOLOv6算法新鲜出炉--训练自己数据集过程_第9张图片
YOLOv6算法新鲜出炉--训练自己数据集过程_第10张图片
yolov6s模型测试效果
YOLOv6算法新鲜出炉--训练自己数据集过程_第11张图片
YOLOv6算法新鲜出炉--训练自己数据集过程_第12张图片
YOLOv6算法新鲜出炉--训练自己数据集过程_第13张图片
YOLOv6算法新鲜出炉--训练自己数据集过程_第14张图片

四、总结

从测试效果来看,yolov6的泛化能力要比yolov5强,虽然,得分稍微比yolov5低点儿,实际工程应用,个人觉得yolov6更好,期待后面部署

你可能感兴趣的:(目标识别检测,深度学习,计算机视觉,算法,深度学习,目标检测,pytorch,人工智能)