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

Posted 2022-07-01 雨果先生

---

前言

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

YOLOv6算法背景：YOLOv6 是美团视觉智能部研发的一款目标检测框架，致力于工业应用。本框架同时专注于检测的精度和推理效率，在工业界常用的尺寸模型中：YOLOv6-nano 在 COCO 上精度可达 35.0% AP，在 T4 上推理速度可达 1242 FPS；YOLOv6-s 在 COCO 上精度可达 43.1% AP，在 T4 上推理速度可达 520 FPS。
评估指标如下;

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

一、环境搭建

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

代码下载后的结构如下图：

安装requirements.txt中的安装包，可以使用清华源快速安装，安装命令：

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

如果出现个别安装不上的，建议单独通过pip install xxx来安装，requirements.txt内容如下：

博主在安装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删除。

路径结尾带有“/”，错误的yaml

配置好自己训练数据集对应的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与YOLOv5效果测试对比

分别用电梯内电单车+人形数据集训练两种算法的模型，模型迭代200次。
YOLOv5模型实际测试效果如下：




yolov6s模型测试效果

四、总结

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