Yolov5 训练

搭建yolo v5 训练环境

yolo v5 本身是一个机器学习图像识别开源包，其中呢，要对其进行训练呢需要准备三个数据集：

• 项目文件夹（project）

  • 训练集（train） ： 用于图片训练

  • images 存放标注图片

  • labels 存放由labelimg标注工具保存的 txt 文件（标签）

  • 验证集(val) ： 用于图片验证

  • images 存放标注图片

  • labels 存放由labelimg标注工具保存的 txt 文件（标签）

  • 测试集(test) ： 用于图片测试

  • images 存放标注图片

  • labels 存放由labelimg标注工具保存的 txt 文件（标签）

由于数据集如果手动分发，过于繁琐而且又不适用于数据集训练

所以以下是 一个数据集自动分发的脚本

# -*- coding: utf-8 -*-
# @Author : dyz

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

# 原始路径
image_original_path = '/home/zqht/workspace/tagging_image/0/'
label_original_path = '/home/zqht/workspace/result_image/0/'
# 训练集路径
train_image_path = '/home/zqht/miniconda3/envs/mytorch/yolov5/traffic_sign/train/images/'
train_label_path = '/home/zqht/miniconda3/envs/mytorch/yolov5/traffic_sign/train/labels/'
# 验证集路径
val_image_path = '/home/zqht/miniconda3/envs/mytorch/yolov5/traffic_sign/val/images/'
val_label_path = '/home/zqht/miniconda3/envs/mytorch/yolov5/traffic_sign/val/labels/'
# 测试集路径
test_image_path = '/home/zqht/miniconda3/envs/mytorch/yolov5/traffic_sign/test/images/'
test_label_path = '/home/zqht/miniconda3/envs/mytorch/yolov5/traffic_sign/test/labels/'

# 数据集划分比例，训练集80%，验证集10%，测试集10%
train_percent = 0.8
val_percent = 0.1
test_percent = 0.1

# 检查文件夹是否存在
def mkdir():
    if not os.path.exists(train_image_path):
        os.makedirs(train_image_path)
    if not os.path.exists(train_label_path):
        os.makedirs(train_label_path)

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

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


def main():
    mkdir()

    total_txt = os.listdir(label_original_path)
    num_txt = len(total_txt)
    list_all_txt = range(num_txt)  # 范围 range(0, num)

    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
    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)
    # 检查两个列表元素是否有重合的元素
    # set_c = set(val_test) & set(val)
    # list_c = list(set_c)
    # print(list_c)
    # print(len(list_c))

    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 + '.jpg'
        srcLabel = label_original_path + name + '.txt'

        if i in train:
            dst_train_Image = train_image_path + name + '.jpg'
            dst_train_Label = train_label_path + name + '.txt'
            shutil.copy(srcImage, dst_train_Image)
            shutil.copy(srcLabel, dst_train_Label)
        elif i in val:
            dst_val_Image = val_image_path + name + '.jpg'
            dst_val_Label = val_label_path + name + '.txt'
            shutil.copy(srcImage, dst_val_Image)
            shutil.copy(srcLabel, dst_val_Label)
        else:
            dst_test_Image = test_image_path + name + '.jpg'
            dst_test_Label = test_label_path + name + '.txt'
            shutil.copy(srcImage, dst_test_Image)
            shutil.copy(srcLabel, dst_test_Label)

    print "Sir Copy function is completed"


if __name__ == '__main__':
    main()

分发完成后 三个训练集文件夹下都存在数据了,就可以修改（没有的可以创建）data文件夹下的mydata.yaml文件了，修改成以下文本即可。

train: ./coco/train/images/# 训练集图片目录
val: ./coco/val/images/#验证集图片目录
test: ./coco/test/images/#测试集图片目录
nc: 2 #number_classes 要识别的类别总数
names: [ 'Curry', 'Durant' ] #此列表里是一共要识别的类别

现在已全部完成环境设置

可以运行train.py 文件 对 数据集进行训练了

python train.py --batch 16 --epochs 300 --data ./data/mydata.yaml --weights ./weights/yolov5s.pt

训练得到的权重文件保存在yolov5/runs/train/文件夹中。

train命令中可使用的参数及其含义如下：
--weights		指定权重，如果不加此参数会默认使用官方预训练的yolov5s.pt
--cfg 		指定模型文件
--data		指定数据文件 
--hyp		指定超参数文件
--epochs		指训练完整数据的次数，默认300
--batch-size	指一次迭代训练的数据大小，默认16，官方推荐越大越好，用你GPU
能承受最大的。可简写为--batch
--img-size 	指定训练图片大小，默认640，可简写为--img
--name 		指定结果文件名，默认result.txt
--device 		指定训练设备，如--device 0,1,2,3
--local_rank 	分布式训练参数，不要自己修改！
--log-imgs 	W&B的图片数量，默认16，最大100
--workers 	指定dataloader的workers数量，默认8
--project 		训练结果存放目录，默认./runs/train/
--name 		训练结果存放名，默认exp

使用训练好的权重文件对图片进行检测

							#图片目录路径					使用训练好的权重文件
python detect.py --source ./coco/images/val/ --weights ./runs/train/exp3/weights/best.pt  --conf 0.5