YOLOv5模型训练流程

此文章只是记录使用，以便后续查看，不作为教程，刚接触，可能有错误

YOLOv5模型训练流程

一、数据集的准备

1.在源码根目录新建mydata文件夹，在此文件夹下新建images和labels文件夹

目录树如下：

├───mydata
│   ├───images
│   └───labels

2.在images下放需要标记的所有图片，然后使用labelimg进行图片标记

3.打开labelimg，点击菜单栏View下的Auto Save Mode使得勾选，点击左边状态栏的PascalVOC切换为YOLO，再点击左边状态栏的Change Save Dir选择路径为上面mydata下的labels文件夹，点击Open Dir开始标记图片。标记完成后剪切出labels文件夹下的classes.txt到mydata目录下

4.根目录新建datasets/defect(随机，但后面路径要一致)文件夹，在此文件夹下新建images和labels文件夹，还需新建的文件

夹目录树如下：

├───datasets
│   └───defect
│       ├───images
│       │   ├───test
│       │   ├───train
│       │   └───val
│       └───labels
│           ├───test
│           ├───train
│           └───val

5.在根目录新建data.py

# 将图片和标注数据按比例切分为 训练集和测试集
import shutil
import random
import os

# 原始路径
image_original_path = "./mydata/images/"
label_original_path = "./mydata/labels/"

cur_path = os.getcwd()

# 训练集路径
train_image_path = os.path.join(cur_path, "datasets/defect/images/train/").replace(os.sep, "/")
train_label_path = os.path.join(cur_path, "datasets/defect/labels/train/").replace(os.sep, "/")

# 验证集路径
val_image_path = os.path.join(cur_path, "datasets/defect/images/val/").replace(os.sep, "/")
val_label_path = os.path.join(cur_path, "datasets/defect/labels/val/").replace(os.sep, "/")

# 测试集路径
test_image_path = os.path.join(cur_path, "datasets/defect/images/test/").replace(os.sep, "/")
test_label_path = os.path.join(cur_path, "datasets/defect/labels/test/").replace(os.sep, "/")

# 训练集目录
list_train = os.path.join(cur_path, "datasets/defect/train.txt").replace(os.sep, "/")
list_val = os.path.join(cur_path, "datasets/defect/val.txt").replace(os.sep, "/")
list_test = os.path.join(cur_path, "datasets/defect/test.txt").replace(os.sep, "/")

train_percent = 0.6
val_percent = 0.2
test_percent = 0.2


def del_file(path):
    for i in os.listdir(path):
        file_data = path + "\\" + i
        os.remove(file_data)


def mkdir():
    if not os.path.exists(train_image_path):
        os.makedirs(train_image_path)
    else:
        del_file(train_image_path)
    if not os.path.exists(train_label_path):
        os.makedirs(train_label_path)
    else:
        del_file(train_label_path)

    if not os.path.exists(val_image_path):
        os.makedirs(val_image_path)
    else:
        del_file(val_image_path)
    if not os.path.exists(val_label_path):
        os.makedirs(val_label_path)
    else:
        del_file(val_label_path)

    if not os.path.exists(test_image_path):
        os.makedirs(test_image_path)
    else:
        del_file(test_image_path)
    if not os.path.exists(test_label_path):
        os.makedirs(test_label_path)
    else:
        del_file(test_label_path)


def clearfile():
    if os.path.exists(list_train):
        os.remove(list_train)
    if os.path.exists(list_val):
        os.remove(list_val)
    if os.path.exists(list_test):
        os.remove(list_test)


def main():
    mkdir()
    clearfile()

    file_train = open(list_train, 'w')
    file_val = open(list_val, 'w')
    file_test = open(list_test, 'w')

    total_txt = os.listdir(label_original_path)
    num_txt = len(total_txt)
    list_all_txt = range(num_txt)

    num_train = int(num_txt * train_percent)
    num_val = int(num_txt * val_percent)
    num_test = num_txt - num_train - num_val

    train = random.sample(list_all_txt, num_train)
    # train从list_all_txt取出num_train个元素
    # 所以list_all_txt列表只剩下了这些元素
    val_test = [i for i in list_all_txt if not i in train]
    # 再从val_test取出num_val个元素，val_test剩下的元素就是test
    val = random.sample(val_test, num_val)

    print("训练集数目：{}, 验证集数目：{}, 测试集数目：{}".format(len(train), len(val), len(val_test) - len(val)))
    for i in list_all_txt:
        name = total_txt[i][:-4]

        srcImage = image_original_path + name + '.png'
        srcLabel = label_original_path + name + ".txt"

        if i in train:
            dst_train_Image = train_image_path + name + '.png'
            dst_train_Label = train_label_path + name + '.txt'
            shutil.copyfile(srcImage, dst_train_Image)
            shutil.copyfile(srcLabel, dst_train_Label)
            file_train.write(dst_train_Image + '\n')
        elif i in val:
            dst_val_Image = val_image_path + name + '.png'
            dst_val_Label = val_label_path + name + '.txt'
            shutil.copyfile(srcImage, dst_val_Image)
            shutil.copyfile(srcLabel, dst_val_Label)
            file_val.write(dst_val_Image + '\n')
        else:
            dst_test_Image = test_image_path + name + '.png'
            dst_test_Label = test_label_path + name + '.txt'
            shutil.copyfile(srcImage, dst_test_Image)
            shutil.copyfile(srcLabel, dst_test_Label)
            file_test.write(dst_test_Image + '\n')

    file_train.close()
    file_val.close()
    file_test.close()


if __name__ == "__main__":
    main()

执行代码。
注意：根据图片的后缀不同，修改代码中.png，根据路径不同修改路径
至此，数据集准备完成。

二、配置训练参数

1.在yolov5-master/data下新建my.yaml

# YOLOv5  by Ultralytics, AGPL-3.0 license
# COCO 2017 dataset http://cocodataset.org by Microsoft
# Example usage: python train.py --data coco.yaml
# parent
# ├── yolov5
# └── datasets
#     └── coco  ← downloads here (20.1 GB)


# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: E:\1_Work\0802\yolov5-master\datasets\defect  # dataset root dir
train: images/train  # train images (relative to 'path') 118287 images
val: images/val  # val images (relative to 'path') 5000 images
test: images/test  # 20288 of 40670 images, submit to https://competitions.codalab.org/competitions/20794

# Classes
nc: 3
names: ["piaochong", "qicao", "changchong"]

注意：修改对应的path,train,val,test路径
nc为标记的类别数量
names为类别名称，在之前的classes.txt中复制

2.在yolov5-master/models下复制yolov5s.yaml文件重命名为yolov5s_test.yaml，修改nc类别数量

# YOLOv5  by Ultralytics, AGPL-3.0 license

# Parameters
nc: 3  # 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)
  ]

3.修改train.py：

修改项	值	解释
--weights	default=ROOT / ‘yolov5s.pt’	预训练权重，在官方Github下载
--cfg	default=‘models/yolov5s_test.yaml’	加载模型，之前复制修改的文件
--data	default=ROOT / ‘data/my.yaml’	配置文件
--epochs	default=2	训练批次
--batch-size	default=2	每批次的输入数据量

修改完成之后运行即可开始训练。

三、检验训练的模型的推理效果

修改detect.py文件：

修改项	值	解释
--weights	default=ROOT / ‘best.pt’	为训练完成后在runs/train/exp/weights下的best.pt文件
--source	default=ROOT / ‘datasets/defect/images/test’	测试数据目录

执行文件即可
结果在yolov5-master/runs/detect/exp下。
注意：如果生成的图片没有框，可以降低置信度排查错误，即detect.py文件中的–conf-thres，–iou-thres。