import os
import shutil

nunber_class = 1                                         # 训练集中的类别数,按需修改!
model_version = 'yolov5s'                                # yolov5模型，按需修改!
my_data_path = 'my_yolo_dataset'                         # 自己符合yolo格式数据集的路径，按需修改!
names = ['person']                                       # 数据集中的标签名字，按需修改!
save_name = 'my_yolov5'                                  # 保存的名字
data_version = 'coco128'
model_pah = 'models'                                     # 模型yaml文件所在的文件夹
data_path = 'data'                                       # data path
save_model_yaml_path = os.path.join(model_pah, save_name + '.yaml')
save_data_yaml_path = os.path.join(data_path, save_name + '.yaml')
model_yaml_path = os.path.join(model_pah, model_version + '.yaml')
data_yaml_path = os.path.join(data_path, data_version + '.yaml')


def Change_model_yaml():

    shutil.copyfile(model_yaml_path, save_model_yaml_path)     # 先复制一份
    with open(save_model_yaml_path, 'r') as f:
        text = f.read().splitlines()
        text[3] = 'nc: ' + str(nunber_class) + ' # number of classes'

    with open(save_model_yaml_path, 'w') as w:
        for cont in text:
            cont = cont + '\n'
            w.write(cont)
    print("model yaml has been done .....")


def change_data_yaml():

    shutil.copyfile(data_yaml_path, save_data_yaml_path)     # 先复制一份
    with open(save_data_yaml_path, 'r', encoding='utf-8') as f:
        text = f.read().splitlines()[11:18]
        text[0] = "train: " + os.path.join(my_data_path, "train")
        text[1] = "val: " + os.path.join(my_data_path, "val")
        text[5] = "nc: " + str(nunber_class)
        text[6] = "names: " + str(names)

    with open(save_data_yaml_path, 'w', encoding='utf-8') as w:
        for cont in text:
            cont = cont + '\n'
            w.write(cont)

    print("data yaml has been done .....")


Change_model_yaml()
change_data_yaml()

运行完会看到多了两个yaml文件

4.修改train.py中的参数配置文件

在train.py下面有一个参数配置文件，我们主要修改以下几个

parser.add_argument('--weights', type=str, default='weights/yolov5s.pt', help='initial weights path')           # 修改成需要的pretrain权重
parser.add_argument('--cfg', type=str, default='models/my_yolov5.yaml', help='model.yaml path')                 # 和权重对应的模型配置文件，由脚本生成
parser.add_argument('--data', type=str, default='data/my_yolov5.yaml', help='dataset.yaml path')                # 自己数据集相关的配置文件，脚本生成
parser.add_argument('--hyp', type=str, default='data/hyps/hyp.scratch.yaml', help='hyperparameters path')       # 超参数
parser.add_argument('--epochs', type=int, default=300)                                                          # 训练epoch
parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs')                 # batchsize根据电脑配置修改

这里说一下，这个权重要和上面脚本中的model_version对应

batch_size如果电脑配置不太好，尽量设置小一点。

还有epochs的设置，如果电脑不太好，可以先设置小一点，跑个效果，yolov5会自动保留这些epoch中最好的结果，可以在runs/train文件夹中找到，下次可以直接在用这个权重继续训练！

设置完后，就可以直接运行train.py进行训练了！

Pytorch机器学习（四）——YOLOV5训练自己的VOC数据集

Pytorch机器学习（四）——YOLOV5训练自己的数据集（以VOC格式为例

前言

一、准备工作

1.下载源码

2.将符合要求的数据集放到根目录下

3.运行训练前需要的配置脚本

4.修改train.py中的参数配置文件

Guess you like