使用yolov5训练自己的数据集（GPU版本）

使用yolov5训练自己的数据集（GPU版本）

一、前言介绍

由于yolov5版本众多，本人用的是最新的7.0版本

本人使用：

• PyTorch==1.12.1
• Python=3.9
• Cuda=11.3

注意：此版本需要Python版本要大于等于3.8.0，PyTorch版本大于等于1.8

二、YOLO源码下载

YOLOv5源码地址

如果要选择其他版本可以在Tags内找到

三、使用conda配置虚拟环境

conda creat -n pytorch python=3.9
conda activate pytorch

创建并激活好虚拟环境时，cd进入YOLOv5目录安装项目依赖

pip install -r requirements.txt comet_ml -i  https://pypi.tuna.tsinghua.edu.cn/simple

可以顺带把comet_ml安装了（本人运行以后发现没装好）若还有什么其他没有装好的依赖，看报错一步步安装即可

四、数据集的准备

打标签可以看我之前写的Labelimg的安装与使用，个人建议打xml格式的，因为网上都是xml转成txt，或者如果是txt则需要转换成xml再转回来（有点脱裤子放p~了）所以就直接打成voc格式就好了

首先创建文件夹，创建的文件结构如下：

My_Data/
├── Annotations # 放xml文件
├── ImageSets # 先不用管
└── JPEGImages # 放数据集

打标签使用labelimg这里不多做介绍，CSDN有很多（使用Anaconda安装）

五、训练集，测试集，验证集的划分

首先在My_Data文件夹下创建一个classes.names的文件

按照标签的类别一行一个放好

将以下代码保存并运行：

import glob
import random
import os
import sys

BASE_PATH = sys.path[0]
print(BASE_PATH)
XML_PATH = os.path.join(BASE_PATH, 'Annotations')
JPG_PATH = os.path.join(BASE_PATH, 'JPEGImages')
TXT_PATH = os.path.join(BASE_PATH, 'ImageSets')

def split(full_list, shuffle=False):
    n_total = len(full_list)
    offset1 = int(n_total * 0.05)
    offset2 = int(n_total * 0.2)
    if n_total == 0:
        return [], full_list
    if shuffle:
        random.shuffle(full_list)

    sublist_1 = full_list[:offset1]
    sublist_2 = full_list[offset1:offset2]
    sublist_3 = full_list[offset2:]
    return sublist_1, sublist_2, sublist_3

def generate_train_and_val():
    file_list = []
    for xml_file in glob.glob(str(os.path.join(XML_PATH,'*.xml'))):
        base_name = os.path.basename(xml_file)
        jpg_file = os.path.join(JPG_PATH, base_name[:-4]+'.jpg')
        if os.path.exists(jpg_file):
            file_list.append(jpg_file)

    file_list = [''.join([x + '\n']) for x in file_list]

    test, val, train = split(file_list, True)
    with open(os.path.join(TXT_PATH, 'all.txt'), 'w') as tf:
        for t in file_list:
            tf.write(t)
    with open(os.path.join(TXT_PATH, 'train.txt'),'w') as tf:
        for t in train:
            tf.write(t)
    with open(os.path.join(TXT_PATH, 'val.txt'),'w') as tf:
        for t in val:
            tf.write(t)
    with open(os.path.join(TXT_PATH, 'test.txt'),'w') as tf:
        for t in test:
            tf.write(t)

generate_train_and_val()

该段代码会扫描Annotations文件夹和JPEGImages文件夹同名的xml和jpg文件，并将其写到all.txt文件中

并且程序会按照all.txt的全部内容按照80% ；15%； 5%分成“train.txt”、“val.txt”、“test.txt”三个文件，保存在“ImageSets”文件夹内。

最后可以看到文件结构目录如下

My_Data/
├── Annotations
├── classes.names
├── ImagSets
├── JPEGImages
├── txt_gen.py
└── xml2yolo.py

我们在利用xml转yolo格式的程序，注意将classes改成自己的类别：

import os
import sys
import glob
import xml.etree.ElementTree as ET


#选择输入路径
BASE_PATH = sys.path[0]
ANNOTATIONS_PATH = os.path.join(BASE_PATH, 'Annotations')
LABELS_PATH = os.path.join(BASE_PATH, 'JPEGImages')

#类名
# classes = ["powdery mildew", "verticillium wilt", "gray mold", "bacterial wilt", "anthracnose", "leaf spot", "red spider", "thrips", "Beet armyworm", "Spodoptera litura"]   # 改成自己的类别
classes = ["powdery mildew"]
with open(os.path.join(BASE_PATH, 'classes.names')) as f:
    class_names = f.readlines()
class_names = [x.strip() for x in class_names]
class_count = [0] * len(class_names)

#转换一个xml文件为txt
def single_xml_to_txt(xml_file):
    tree = ET.parse(xml_file)
    root = tree.getroot()
    #保存txt文件路径
    txt_file = os.path.join(LABELS_PATH, os.path.basename(xml_file)[:-4] + '.txt')
    with open(txt_file, 'w') as tf:
        for member in root.findall('object'):
	        #从xml获取图像的宽和高
            picture_width = int(root.find('size')[0].text)
            picture_height = int(root.find('size')[1].text)
            class_name = member[0].text

            #类名对应的index
            class_num = class_names.index(class_name)
            class_count[class_num] += 1
            box_x_min = int(member[4][0].text)  # 左上角横坐标
            box_y_min = int(member[4][1].text)  # 左上角纵坐标
            box_x_max = int(member[4][2].text)  # 右下角横坐标
            box_y_max = int(member[4][3].text)  # 右下角纵坐标

            # 转成相对位置和宽高（所有值处于0~1之间）
            x_center = (box_x_min + box_x_max) / (2 * picture_width)
            y_center = (box_y_min + box_y_max) / (2 * picture_height)
            width = (box_x_max - box_x_min) / picture_width
            height = (box_y_max - box_y_min) / picture_height
            #print(class_num, x_center, y_center, width, height)
            tf.write(str(class_num) + ' ' + str(x_center) + ' ' + str(y_center) + ' ' + str(width) + ' ' + str(
                height) + '\n')


#
with open('classes.names', 'w') as f:
    for c in class_names:
        f.write(c + '\n')

#  转换文件夹下的所有xml文件为txt
for xml_file in glob.glob(str(os.path.join(ANNOTATIONS_PATH, '*.xml'))):
    #print(xml_file)
    try:
        single_xml_to_txt(xml_file)
    except Exception as e:
        print(e, xml_file)


for c, n in zip(class_names, class_count):
    print(c + ':' + str(n))

程序运行以后，可以看到JPEGImg文件夹内放着对应的txt文件

六、配置文件

打开yolov5文件夹，在data目录下创建一个my_data.yaml的文件

内容如下：

train: /home/liu/Desktop/My_Data/ImageSets/train.txt
val: /home/liu/Desktop/My_Data/ImageSets/val.txt

nc: 1

names: ["powdery mildew"]

注意冒号之后要有空格，names的顺序要和classes.names中一致

去官网上下载预训练模型，我使用的是yolov5s.pt

我们进入modle目录，找到yolov5s.yaml文件复制并改名

把nc修改成你自己的种类数

七、开始训练

我们可以打开train.py文件，找到训练模型需要的一些参数

if __name__ == '__main__':
"""
    opt模型主要参数解析：
    --weights：初始化的权重文件的路径地址
    --cfg：模型yaml文件的路径地址
    --data：数据yaml文件的路径地址
    --hyp：超参数文件路径地址
    --epochs：训练轮次
    --batch-size：喂入批次文件的多少
    --img-size：输入图片尺寸
    --rect:是否采用矩形训练，默认False
    --resume:接着打断训练上次的结果接着训练
    --nosave:不保存模型，默认False
    --notest:不进行test，默认False
    --noautoanchor:不自动调整anchor，默认False
    --evolve:是否进行超参数进化，默认False
    --bucket:谷歌云盘bucket，一般不会用到
    --cache-images:是否提前缓存图片到内存，以加快训练速度，默认False
    --image-weights：使用加权图像选择进行训练
    --device:训练的设备，cpu；0(表示一个gpu设备cuda:0)；0,1,2,3(多个gpu设备)
    --multi-scale:是否进行多尺度训练，默认False
    --single-cls:数据集是否只有一个类别，默认False
    --adam:是否使用adam优化器
    --sync-bn:是否使用跨卡同步BN,在DDP模式使用
    --local_rank：DDP参数，请勿修改
    --workers：最大工作核心数
    --project:训练模型的保存位置
    --name：模型保存的目录名称
    --exist-ok：模型目录是否存在，不存在就创建
"""
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', type=str, default='yolov5s.pt', help='initial weights path')
    parser.add_argument('--cfg', type=str, default='', help='model.yaml path')
    parser.add_argument('--data', type=str, default='data/coco128.yaml', help='data.yaml path')
    parser.add_argument('--hyp', type=str, default='data/hyp.scratch.yaml', help='hyperparameters path')
    parser.add_argument('--epochs', type=int, default=300)
    parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs')
    parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='[train, test] image sizes')
    parser.add_argument('--rect', action='store_true', help='rectangular training')
    parser.add_argument('--resume', nargs='?', const=True, default=False, help='resume most recent training')
    parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
    parser.add_argument('--notest', action='store_true', help='only test final epoch')
    parser.add_argument('--noautoanchor', action='store_true', help='disable autoanchor check')
    parser.add_argument('--evolve', action='store_true', help='evolve hyperparameters')
    parser.add_argument('--bucket', type=str, default='', help='gsutil bucket')
    parser.add_argument('--cache-images', action='store_true', help='cache images for faster training')
    parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training')
    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%%')
    parser.add_argument('--single-cls', action='store_true', help='train multi-class data as single-class')
    parser.add_argument('--adam', action='store_true', help='use torch.optim.Adam() optimizer')
    parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode')
    parser.add_argument('--local_rank', type=int, default=-1, help='DDP parameter, do not modify')
    parser.add_argument('--workers', type=int, default=8, help='maximum number of dataloader workers')
    parser.add_argument('--project', default='runs/train', help='save to project/name')
    parser.add_argument('--entity', default=None, help='W&B entity')
    parser.add_argument('--name', default='exp', help='save to project/name')
    parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
    parser.add_argument('--quad', action='store_true', help='quad dataloader')
    parser.add_argument('--linear-lr', action='store_true', help='linear LR')
    parser.add_argument('--label-smoothing', type=float, default=0.0, help='Label smoothing epsilon')
    parser.add_argument('--upload_dataset', action='store_true', help='Upload dataset as W&B artifact table')
    parser.add_argument('--bbox_interval', type=int, default=-1, help='Set bounding-box image logging interval for W&B')
    parser.add_argument('--save_period', type=int, default=-1, help='Log model after every "save_period" epoch')
    parser.add_argument('--artifact_alias', type=str, default="latest", help='version of dataset artifact to be used')
    opt = parser.parse_args()

我们只需要用一些常用的开始训练即可

# 激活conda环境
conda activate pytorch
# 开始训练
python3 train.py --img-size 640 --weights yolov5s.pt --data data/My_data.yaml --cfg model/yolov5s.yaml --batch-size 4 --epochs 100

其中batch-size和epochs如果显卡不行你就调小一点

被训练好的模型会被放在runs/train/expx/weights目录下

八、训练结束

本次的教程并没有结束，由于时间问题，我就先写这么多