YOLOV5训练自己数据集

2023-11-04 13:53:40

系列文章目录

使用 YOLOV5 训练自己数据集例程

前言

【个人学习笔记记录，如有错误，欢迎指正】

YOLO-V5 代码仓库地址：https://github.com/ultralytics/yolov5

一、环境配置

这里在 Windows 下跑的代码！！！
这里不介绍环境配置，如果需要文档，请留言。

二、需注意的配置参数介绍

在 train 脚本中，找到 parse_opt 函数。

parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s.pt', help='initial weights path')
parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path')
parser.add_argument('--epochs', type=int, default=300)
parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs, -1 for autobatch')
parser.add_argument('--workers', type=int, default=8, help='max dataloader workers (per RANK in DDP mode)')

1. --weights

这个是模型需要的权值文件，可以自己下载对应的版本，放在 yolov5-master 文件夹下，我这里默认使用 yolov5s.pt。
注意：1：这里 yolo5s.pt 和 yolo5l.pt 是对应不同的模型大小和深度；
2：每一个权值文件都分别对应一个 yaml 文件。

2. --data

自己数据集的 yaml 文件的路径，在下一步会教怎么实现自己数据集的yaml文件。

3. --epochs

迭代整个训练数据集的次数，默认 300，可自行修改。

4. --batch-size

每个训练批次中图像的个数，默认 16，对于显存不够的同学，建议改小一点。例如：8，4等。

5. --workers

进程数量，默认为 8，对于 Windows 下需要将此数值改为 0。

三、yaml 的修改

1.修改数据集路径

在修改 yaml 文件之前，需要对自己数据集的文件夹格式进行修改。
训练图像和标签在同一级文件目录下

|——data
	|-- NEU-DET
			|-- train
				 |-- images
				 |-- labels
			|-- valid
				 |-- images
				 |-- labels

2.修改 yaml

这里给出 coco.yaml 文件的内容。

创建一个 data.yaml 文件，建议存放在数据集文集根目录下。

根据第一步的数据集的路径格式，再参考coco.yaml 文件修改为自己的数据集：

【train ：‘path’】训练数据集的路径
【val：‘path’】训练数据集的路径
【nc：‘number of classes’】需要分类的类别个数
【names：‘names’】每个类别对应的类别名称【注：这个的排序和标注文件的排序是一样的！】

这里 “:” 冒号之前的内容不能修改！且都是英文输入法，。

3. 配置 opt

将 parse_opt 函数中的 default 修改为自己的创建的 data.yaml 文件路径

parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path')

需要将 “ROOT /” 删除，然后修改为

parser.add_argument('--data', type=str, default= 'G:\\data\\NEU-DET\\data.yaml', help='dataset.yaml path')

这样就可以开始训练自己的数据集了。

总结

基于 YOLO-V5 训练自己的数据集，整体来说就这几步，还算简单。

码农公寓