文章目录
前言
本文主要介绍如何利用YOLOv5训练自己的数据集
1、数据标注说明
以生活垃圾数据集为例子
生活垃圾数据集组成：
python
【有害垃圾】：电池（1 号、2 号、5 号）、过期药品或内包装等；
【可回收垃圾】：易拉罐、小号矿泉水瓶；
【厨余垃圾】：小土豆、切过的白萝卜、胡萝卜，尺寸为电池大小；
【其他垃圾】：瓷片、鹅卵石（小土豆大小）、砖块等。YOLO数据有三个要点
- images，存放图片
- labes，对应Images图片的标签
- data_txt, 划分images图片的数据集，形成三个txt
文件结构：
python
----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
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
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 names3、训练模型
1、激活python环境
python
conda activate reid-pytorch # 这里我是激活reid-pytorch环境2、在服务器上训练自己的模型
python
# 注意，这里是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 &
python
# 如果是本地电脑上训练，直接在pycharm的train.py添加上述参数，直接训练
--batch-size 16 --epochs 200
--cfg models/yolov5s.yaml
--data data/waste-classification.yaml
--weights weights/yolov5s.pt
--device 0 # 采用显卡0进行训练
查看自己电脑显卡情况
