前言
本文主要介绍如何利用YOLOv5训练自己的数据集
1、数据标注说明
以生活垃圾数据集为例子
生活垃圾数据集组成:
【有害垃圾】:电池(1 号、2 号、5 号)、过期药品或内包装等;
【可回收垃圾】:易拉罐、小号矿泉水瓶;
【厨余垃圾】:小土豆、切过的白萝卜、胡萝卜,尺寸为电池大小;
【其他垃圾】:瓷片、鹅卵石(小土豆大小)、砖块等。
YOLO数据有三个要点
- images,存放图片
- labes,对应Images图片的标签
- data_txt, 划分images图片的数据集,形成三个txt
文件结构:
----data\
|----classes.txt # 标签种类
|----data-txt\ # 数据集文件集合
| |----test.txt
| |----train.txt
| |----val.txt
|----images\ # 数据集图片
| |----test\
| | |----fimg_23.jpg
| | |----fimg_38.jpg
| | |----.....
| |----train\
| | |----fimg_1.jpg
| | |----fimg_2.jpg
| | |----.....
| |----val\
| | |----fimg_4.jpg
| | |----fimg_6.jpg
| | |----.....
|----labels\ # yolo标签
| |----test\
| | |----fimg_23.txt
| | |----fimg_38.txt
| | |----.....
| |----train\
| | |----fimg_1.txt
| | |----fimg_2.txt
| | |----.....
| |----val\
| | |----fimg_4.txt
| | |----fimg_6.txt
| | |----.....
|----waste-classification.yaml # 数据集配置文件
2、定义自己模型文件
1、定义yolov5网络模型文件
如 models/yolov5l.yaml
# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
# Parameters
nc: 4 # number of classes 类别,这里的类别一共四种
depth_multiple: 0.33 # model depth multiple ,模型宽度
width_multiple: 0.50 # layer channel multiple ,通道数
anchors:
- [10,13, 16,30, 33,23] # P3/8
- [30,61, 62,45, 59,119] # P4/16
- [116,90, 156,198, 373,326] # P5/32
# YOLOv5 v6.0 backbone
backbone:
# [from, number, module, args]
[[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
[-1, 1, Conv, [128, 3, 2]], # 1-P2/4
[-1, 3, C3, [128]],
[-1, 1, Conv, [256, 3, 2]], # 3-P3/8
[-1, 6, C3, [256]],
[-1, 1, Conv, [512, 3, 2]], # 5-P4/16
[-1, 9, C3, [512]],
[-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
[-1, 3, C3, [1024]],
[-1, 1, SPPF, [1024, 5]], # 9
]
# YOLOv5 v6.0 head
head:
[[-1, 1, Conv, [512, 1, 1]],
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
[[-1, 6], 1, Concat, [1]], # cat backbone P4
[-1, 3, C3, [512, False]], # 13
[-1, 1, Conv, [256, 1, 1]],
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
[[-1, 4], 1, Concat, [1]], # cat backbone P3
[-1, 3, C3, [256, False]], # 17 (P3/8-small)
[-1, 1, Conv, [256, 3, 2]],
[[-1, 14], 1, Concat, [1]], # cat head P4
[-1, 3, C3, [512, False]], # 20 (P4/16-medium)
[-1, 1, Conv, [512, 3, 2]],
[[-1, 10], 1, Concat, [1]], # cat head P5
[-1, 3, C3, [1024, False]], # 23 (P5/32-large)
[[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
]
2、在data/目标下,定义自己数据集的配置文件
如data/waste-classification.yaml
train: data/data-txt/train.txt
val: data/data-txt/val.txt
test: data/data-txt/test.txt
# Classes
nc: 4 # number of classes
names: ['recyclable waste', 'hazardous waste','kitchen waste','other waste'] # class names
3、训练模型
1、激活python环境
conda activate reid-pytorch # 这里我是激活reid-pytorch环境
2、在服务器上训练自己的模型
# 注意,这里是linux服务器上训练命令
# nohup 是后台运行,--batch-size 每次送入的图片数, --epochs 训练多少轮 --cfg 模型配置 --data 数据集配置
# --weights 以哪个预训练权重训练 --device 在那个显卡上训练(这里一共有四张显卡) > 重定向,将训练的结果输出到myout.file,方便查看
nohup python train.py --batch-size 16 --epochs 200 --cfg models/yolov5s.yaml --data data/waste-classification.yaml --weights weights/yolov5s.pt --device 0,1,2,3 > myout.file 2>&1 &
# 如果是本地电脑上训练,直接在pycharm的train.py添加上述参数,直接训练
--batch-size 16 --epochs 200
--cfg models/yolov5s.yaml
--data data/waste-classification.yaml
--weights weights/yolov5s.pt
--device 0 # 采用显卡0进行训练
查看自己电脑显卡情况
4、关于预训练模型
如果你的数据集比较大,建议从头开始训练,这可能有助于模型更好地学习数据集的特征。 反之,如果你的数据集较小或者是中等大小,使用预训练模型作为起点可能会得到更好的结果,因为预训练模型已经学习了大量的图像特征,可以提升模型对新样本的泛化能力。
4.1、采用预训练模式
当自身数据量较小的时候,采用预训练模型去训练,即添加 --weight 参数
python train.py --data custom.yaml --weights yolov5s.pt
这个命令是用于训练一个YOLOv5模型。参数--data
后面跟的是一个YAML文件(custom.yaml
),这个文件包含了数据集的配置信息。参数--weights
后面跟的是预训练模型的权重文件,这里用的是yolov5s.pt
,其中yolov5s
代表的是YOLOv5的一个小型版本,适用于在资源有限的设备上进行快速训练和检测。
4.2、无预训练模式
当自身数据量很大的时候,建议采用从零开始训练自己的模型,即不添加 --weight参数
python train.py --data custom.yaml --weights '' --cfg yolov5s.yaml
这个命令同样是用于训练YOLOv5模型,但与上面的命令不同的是,这里--weights
参数后面是空的,意味着训练将从头开始,不使用任何预训练的权重。--cfg
参数后面跟的是模型的配置文件,这里使用的是yolov5s.yaml
,该文件定义了模型的架构和其他参数。
参考
文献