Yolov5+Pytorch训练自己的数据集

最近,我让介个人学习神经网络,但是发现自己也不会。连自己都不会,又怎么帮别人解决问题呢?于是我花了一上午研究了一下 Darknet 训练 Yolov4-tiny(我们上个赛季用的就是这个,所以基本上不用怎么学), 一下午研究了一下 Pytorch 怎么训练 yolov5 模型,故作此篇。

介个人最近在学 yolo,前不久也被环境折磨了好久(可以看这里,在下亲眼见证现状的惨烈哈哈哈哈~),于是我努力学习如何使用 Pytorch ,这样就可以继续嘲笑了哈哈哈哈哈哈哈哈哈~~

Pytorch 训练

参考博客:这里

1. 训练数据集制作

1.1 将图片和标签导入

首先建立一个自己的数据文件夹 armor_coco,并在该文件夹下创建 all_images 和 all_xml 两个文件夹。目录结构如下:

armor_coco
├── all_images
├── all_xml

之后将需要训练的图片和 xml 标签分别存放在 all_images 和 all_xml 两个文件夹下

1.2 可选项:导入已有的 txt 标签

如果此时手上已经有转换好的 txt 文件,则在 armor_coco 文件夹下创建 all_labels 文件夹,将 txt 标签放进去。之后在生成的时候,将 py 脚本中的一行注释掉即可(之后会讲)

1.3 创建 make_txt.py 并执行

make_txt.py 是用来划分数据集使用。

在 armor_coco 目录下创建 make_txt.py,并写入以下内容:

import os
import random
trainval_percent = 0.1  
train_percent = 0.9    
xmlfilepath = 'all_images'
txtsavepath = 'ImageSets'
total_xml = os.listdir(xmlfilepath)
num = len(total_xml)
list = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
trainval = random.sample(list, tv) #从所有list中返回tv个数量的项目
train = random.sample(trainval, tr)
if not os.path.exists('ImageSets/'):
    os.makedirs('ImageSets/')
ftrainval = open('ImageSets/trainval.txt', 'w')
ftest = open('ImageSets/test.txt', 'w')
ftrain = open('ImageSets/train.txt', 'w')
fval = open('ImageSets/val.txt', 'w')
for i in list:
    name = total_xml[i][:-4] + '\n'
    if i in trainval:
        ftrainval.write(name)
        if i in train:
            ftest.write(name)
        else:
            fval.write(name)
    else:
        ftrain.write(name)
ftrainval.close()
ftrain.close()
fval.close()
ftest.close()

运行:

python make_txt.py

运行完成后生成 ImageSets 文件夹,在该文件夹下生成 4 个.txt文件:

├── test.txt
├── train.txt
├── trainval.txt
└── val.txt

这些文件保存着图片的划分信息,划分为 train val test 三类。

1.4 创建 train_val.py 文件并执行

在 armor_coco 文件夹下创建 train_val.py 文件,并写入以下内容:

import xml.etree.ElementTree as ET
import pickle
import os
import shutil
from os import listdir, getcwd
from os.path import join
sets = ['train', 'trainval']

### 需要修改以下内容
classes =  ['Dark1','Dark2','Dark3','Dark4','Dark5','Dark6','Dark7','Red0','Red1','Red2','Red3','Red4','Red5','Red6','Red7','Red8','Blue0','Blue1','Blue2','Blue3','Blue4','Blue5','Blue6','Blue7','Blue8']
def convert(size, box):
    dw = 1. / size[0]
    dh = 1. / size[1]
    x = (box[0] + box[1]) / 2.0
    y = (box[2] + box[3]) / 2.0
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return (x, y, w, h)
def convert_annotation(image_id):
    in_file = open('all_xml/%s.xml' % (image_id))
    out_file = open('all_labels/%s.txt' % (image_id), 'w')
    tree = ET.parse(in_file)
    root = tree.getroot()
    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)
    for obj in root.iter('object'):
        difficult = obj.find('difficult').text
        cls = obj.find('name').text
        if cls not in classes or int(difficult) == 1:
            continue
        cls_id = classes.index(cls)
        xmlbox = obj.find('bndbox')
        b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),
             float(xmlbox.find('ymax').text))
        bb = convert((w, h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
wd = getcwd()
print(wd)
for image_set in sets:
    if not os.path.exists('all_labels/'):
        os.makedirs('all_labels/')
    image_ids = open('ImageSets/%s.txt' % (image_set)).read().strip().split()
    image_list_file = open('images_%s.txt' % (image_set), 'w')
    labels_list_file=open('labels_%s.txt'%(image_set),'w')
    for image_id in image_ids:
        image_list_file.write('%s.jpg\n' % (image_id))
        labels_list_file.write('%s.txt\n'%(image_id))
        convert_annotation(image_id) #如果标签已经是txt格式,将此行注释掉,所有的txt存放到all_labels文件夹。
    image_list_file.close()
    labels_list_file.close()


def copy_file(new_path,path_txt,search_path):#参数1:存放新文件的位置  参数2:为上一步建立好的train,val训练数据的路径txt文件  参数3:为搜索的文件位置
    if not os.path.exists(new_path):
        os.makedirs(new_path)
    with open(path_txt, 'r') as lines:
        filenames_to_copy = set(line.rstrip() for line in lines)
        # print('filenames_to_copy:',filenames_to_copy)
        # print(len(filenames_to_copy))
    for root, _, filenames in os.walk(search_path):
        # print('root',root)
        # print(_)
        # print(filenames)
        for filename in filenames:
            if filename in filenames_to_copy:
                shutil.copy(os.path.join(root, filename), new_path)

#按照划分好的训练文件的路径搜索目标,并将其复制到yolo格式下的新路径
copy_file('./images/train/','./images_train.txt','./all_images')
copy_file('./images/val/','./images_trainval.txt','./all_images')
copy_file('./labels/train/','./labels_train.txt','./all_labels')
copy_file('./labels/val/','./labels_trainval.txt','./all_labels')

注意需要将 classes 数组修改为我们打的标签,并按照一定顺序进行存放(这个顺序很重要,之后会有类似的修改,也要保证顺序一致。)

另外,如果用已经转换好的标签进行训练,则需找到 convert_annotation 一行注释掉,同时需保证 classes 中的标签顺序与之前转换 txt 时的标签书序一致。

执行:

python3 train_val.py

运行结束后 armor_coco 文件夹下的内容:

├── all_images
├── all_labels
├── all_xml
├── images
│   ├── train
│   └── val
├── ImageSets
└── labels
    ├── train
    └── val

至此,数据集制作完成。

2. 训练

2.1 下载 yolov5 的 Pytorch 框架

下载地址:https://github.com/ultralytics/yolov5/tree/v6.0

解压后打开文件夹

2.2 创建 armor_coco.yaml

在上述文件夹中,打开 data 文件夹,复制一份 coco128.yaml 文件并重命名为 armor_coco.yaml。修改以下内容:

  • path: 修改为 armor_coco 文件夹的路径
  • train: 修改为 images/train
  • test: 修改为 images/test
  • nc: 类别数。因为我们现在有 25 个类,修改为 25
  • names: 修改为与上述 classes 一致

修改后内容如下:

# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
# COCO128 dataset https://www.kaggle.com/ultralytics/coco128 (first 128 images from COCO train2017)
# Example usage: python train.py --data coco128.yaml
# parent
# ├── yolov5
# └── datasets
#     └── coco128  ← downloads here


# 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: ../armor_coco  # dataset root dir
train: images/train  # train images (relative to 'path') 128 images
val: images/val  # val images (relative to 'path') 128 images
test:  # test images (optional)

# Classes
nc: 25  # number of classes
names: ['Dark1','Dark2','Dark3','Dark4','Dark5','Dark6','Dark7','Red0','Red1','Red2','Red3','Red4','Red5','Red6','Red7','Red8','Blue0','Blue1','Blue2','Blue3','Blue4','Blue5','Blue6','Blue7','Blue8']


# Download script/URL (optional)
download: https://github.com/ultralytics/yolov5/releases/download/v1.0/coco128.zip

2.3 开始训练

回到 yolov5-6.0 文件目录下,执行:

python3 train.py --data data/armor_coco.yaml

即可开始以默认超参数进行训练

但是默认参数对于我们来说并不一定是最好的,因此我们需要指定一些训练参数。

以下代码段是从 train.py 文件中截取的,是一些我们可以指定配置参数。

def parse_opt(known=False):
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s.pt', help='initial weights path')
    parser.add_argument('--cfg', type=str, default='', help='model.yaml path')
    parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path')
    parser.add_argument('--hyp', type=str, default=ROOT / 'data/hyps/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, -1 for autobatch')
    parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='train, val image size (pixels)')
    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('--noval', action='store_true', help='only validate final epoch')
    parser.add_argument('--noautoanchor', action='store_true', help='disable autoanchor check')
    parser.add_argument('--evolve', type=int, nargs='?', const=300, help='evolve hyperparameters for x generations')
    parser.add_argument('--bucket', type=str, default='', help='gsutil bucket')
    parser.add_argument('--cache', type=str, nargs='?', const='ram', help='--cache images in "ram" (default) or "disk"')
    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('--workers', type=int, default=8, help='maximum number of dataloader workers')
    parser.add_argument('--project', default=ROOT / 'runs/train', help='save to project/name')
    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('--patience', type=int, default=100, help='EarlyStopping patience (epochs without improvement)')
    parser.add_argument('--freeze', type=int, default=0, help='Number of layers to freeze. backbone=10, all=24')
    parser.add_argument('--save-period', type=int, default=-1, help='Save checkpoint every x epochs (disabled if < 1)')
    parser.add_argument('--local_rank', type=int, default=-1, help='DDP parameter, do not modify')

    # Weights & Biases arguments
    parser.add_argument('--entity', default=None, help='W&B: Entity')
    parser.add_argument('--upload_dataset', action='store_true', help='W&B: Upload dataset as artifact table')
    parser.add_argument('--bbox_interval', type=int, default=-1, help='W&B: Set bounding-box image logging interval')
    parser.add_argument('--artifact_alias', type=str, default='latest', help='W&B: Version of dataset artifact to use')

    opt = parser.parse_known_args()[0] if known else parser.parse_args()
    return opt

我们可能需要修改的配置参数如下:

  • weights: 用于指定开始训练所使用的的模型。一般是断点续传时使用,选择之前训练过的模型接着训练。如果不指定,则会使用一个预训练模型进行训练。
  • data: 指定训练时所使用的配置文件,在这里指定我们之前编辑的 armor_coco.yaml 文件,用于告诉程序我们的训练数据集在哪里。
  • epochs: 训练次数。次数越多,模型可能越准确,也可能会发生过拟合,且训练时间越长,一般指定100-200。如果由于训练次数不足导致训练效果不佳,可以使用断点续传继续训练。
  • batch-size: 一次放入多少张图片进行训练,数值最好为 2^n 形式。数值越大,训练效果越好,但是对显存的需求也越大。
  • imgsz:图片大小,数值为32的整数倍。该数字表示将图片以多大的尺寸进行训练,并且之后使用模型识别的时候输入图片的大小是多少。数值越大,则训练和识别效果越好,但对显卡算力的要求也越高。

在这一次的训练中,我选择训练 100 个 epoch 以 480x480 的图片大小进行训练,一次性训练 64 张图片。执行以下命令:

python3 train.py --data data/armor_coco.yaml --epoch 100 --imgsz 480 --batch-size 64

如果环境配置正常,会看到以下输出,说明正在训练:

# 前面省略一大堆看似没什么用的输出...

Starting training for 100 epochs...

     Epoch   gpu_mem       box       obj       cls    labels  img_size
      0/99     8.03G    0.0643   0.01652   0.06404       182       480: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 573/573 [02:41<00:00,  3.54it/s]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 32/32 [00:12<00:00,  2.60it/s]
                 all       4074       7862      0.687      0.303       0.29      0.143

     Epoch   gpu_mem       box       obj       cls    labels  img_size
      1/99     10.4G   0.04388   0.01098   0.02733       171       480: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 573/573 [02:37<00:00,  3.64it/s]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 32/32 [00:11<00:00,  2.68it/s]
                 all       4074       7862      0.614       0.58      0.572      0.319

     Epoch   gpu_mem       box       obj       cls    labels  img_size
      2/99     10.4G   0.03726   0.01031   0.01695       187       480: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 573/573 [02:38<00:00,  3.63it/s]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 32/32 [00:12<00:00,  2.60it/s]
                 all       4074       7862      0.612      0.696      0.723      0.427

     Epoch   gpu_mem       box       obj       cls    labels  img_size
      3/99     10.4G   0.03242  0.009914   0.01221       178       480: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 573/573 [02:40<00:00,  3.58it/s]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 32/32 [00:12<00:00,  2.61it/s]
                 all       4074       7862      0.753       0.83      0.858      0.508

# 后面省略一大堆在未来会输出的内容~

我们设置的一个 batch 是 64 张图片,一个 epoch 中训练了 573 次。 573*64=36672 跟我们的训练集的图片数量差不多。因此,一个 epoch 相当于是把整个训练集的图片都拿来训练了一遍。

训练结束后,可以在 runs 文件夹下看到训练好的模型(.pt)和验证结果。之后就可以愉悦地使用模型进行识别啦~

  • 1
    点赞
  • 4
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值