使用yolov5训练自己的目标检测数据集

yolov4出来后不久，又出现了yolov5，没有论文。虽然作者没有放上和yolov4的直接测试对比，但在COCO数据集的测试效果还是很可观的。很多人考虑到YOLOv5的创新性不足，对算法是否能够进化，因此yolov5的名头仍有争议。但是既然github上有如此多的star，说明还是受到大多数人认可的一个工程。

并且yolov5是十分容易上手使用的一个目标检测开源项目，只需要改几个文件，配置好自己的数据集格式和目录结构，即可开始训练，对新手十分友好。可以先把训练跑起来，在根据需要去细扣其中的模块设计。

并且，yolov5项目一直在更新维护，除了日常处理bug，还有更优算法的版本（如激活函数的更新）更迭和新的训练可视化工具wandb的加入等，这也是十分难得的。

今天，笔者就来介绍一下自己第一次使用yolov5训练自己的数据集的过程。

工程项目准备

PyTorch等其他依赖的安装这里就不说了，网上有很多教程。

首先从yolov5 github上克隆下代码：

git clone https://github.com/ultralytics/yolov5.git

然后从github下载预训练权重(.pt文件)，没有百度云版本，这里笔者提供下载好的权重文件放到云盘（密码: afp3）供不方便的朋友直接下载。下载好后放到yolov5/weights文件夹下。

数据集准备

yolov5项目所需的数据集文件路径格式如下：

---yolov5
---yolo_data---train---images---labels---val---images---labels---data.yaml

数据集文件夹放在与工程文件同级的目录下。其中包含两个子文件夹和一个配置文件data.yaml，两个子文件夹分别存放训练样本和验证样本，注意images和labels文件夹的名字不要改。

先说一下data.yaml配置文件内的内容

train: ../yolo_data/train
val: ../yolo_data/valnc: [你的检测数据集的类别数]
names: [你的检测数据集的类别名称列表]

共有四个键值对，train，val，nc，names，分别对应如下表：

键	值
train	你的训练数据的文件夹路径
val	你的验证数据的文件夹路径
nc	你的检测数据集的类别数（如COCO数据集中nc=80）
names	你的检测数据集的类别名称列表（如：[‘plane’, ‘car’, ‘dog’, …]）

然后说一下数据集格式的要求：

在train/val目录下要有images和labels两个目录，分别存放图像和标签，其中的文件名要相互对应，分别为.jpg格式和.txt格式，即：

---train---imagesaaaa.jpg    ...---labelsaaaa.txt...

其中图像不必多说，放入你的数据集图像即可，而标签label目录下的txt文件要求格式如下：

类别索引           x           y             w               h
0           0.53857421875 0.8779296875 0.158203125 0.06184895833333333

第一列是图像中每个object的类别索引，注意这里要与data.yaml配置文件中的names相互对应。后面四个列就是该object的检测框坐标了，这里采用的是经过归一化的中心长宽xywh格式，如果你的数据集原生不是这种格式（如对角坐标xyxy）的话，请注意调整。

开始训练

直接运行yolov5/train.py即可，注意配置好参数

python train.py \
--data ../yolo_data/data.yaml \
--cfg models/yolov5s.yaml \
--weights weights/yolov5s.yaml \
--batch-size 8 \
--epochs 100

yolov5最近一次更新增加了wandb（weights and biases，其官网，直接访问不太稳定，建议科学上网访问）训练可视化工具，没有安装的话需要先安装一下，直接pip安装即可

pip install wandb

第一次使用的话会需要注册账号，访问站点https://wandb.ai/authorize，拿到API Key，复制到终端中即可，首次绑定API Key账号后再训练就不在需要输入API Key了，直接将训练可视化结果送到你的wandb账号，可在该站点登录查看。

不仅有训练过程可视化，loss曲线等

甚至还有训练过程中系统的信息的展示：

还是非常简单实用的一个训练可视化工具。

训练完成后再yolov5/runs/train/exp文件夹下可查看一些结果PR曲线，F1曲线等，还有一些样本可视化的结果。

可以看到效果还是可以的（我这个数据集只有猪仔一个类 ^^可爱小猪仔）。

测试

训练完成后可以运行yolov5.detect.py文件，指定一个文件夹，测试其中的图像。

python detect.py \
--weights runs/train/exp/weights/best.pt \
--source ../images/test \
--save-txt

如果设置打开–save-txt，则会保存结果文件。

测试结果会保存在yolov5/runs/detect内。txt结果文件则会保存在yolov5/runs/detect/labels内。

使用yolov5训练自己的数据集就先介绍到这里，如果有问题欢迎留言讨论。