本文主要介绍如何利用YOLOv5训练自己的数据集
以生活垃圾数据集为例子
生活垃圾数据集组成:
【有害垃圾】:电池(1 号、2 号、5 号)、过期药品或内包装等;
【可回收垃圾】:易拉罐、小号矿泉水瓶;
【厨余垃圾】:小土豆、切过的白萝卜、胡萝卜,尺寸为电池大小;
【其他垃圾】:瓷片、鹅卵石(小土豆大小)、砖块等。
YOLO数据有三个要点
文件结构:
----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 # 数据集配置文件
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
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进行训练
查看自己电脑显卡情况