YOLOv5训练自己的数据集【较全面】

前言

这里主要介绍 ultralytics/yolov5 即所谓的 U版YOLO v5，关于YOLO5的原理可以自行搜索相关知识，根据笔者自己测试，效果确实不错，所以这里记录下自己的调试经验。

实验环境

单块 Nvidia GTX 2070 Super, Torch 1.5.0 / torchvision 0.6.0 / python 3.7

代码调试部分

1 源码获取

目前YOLO5社区活跃，更新很快，本人由于不想升级自己的Torch 版本，所以使的是 U版 YOLO 5的 version 2.0
版本，获取链接点击进入
【注意也要下载对应的预训练模型，在网页上的 asserts 部分】

2 自定义数据集

这个其实在官网上也有，点击进入，这里需要注意的是YOLO 5的类别要求是从 0 索引开始的，一定要注意这个，具体的目标图像和标注的目录结构如图所示：

其中标签的格式为 类别索引，实例中心 x y 和 实例宽 高 w h 其中 x,y,w,h均为相对于图像的比值，一定要注意。

2.1 自定义数据集生成标签代码

其实上面的 images中的 train 和 val 的文件夹，如果不想把自己的图像再次复制一遍占用磁盘空间，可以直接使用 软连接
ln -s /yourData Dir /images/train
这样也是可以的，就可以省去复制图像的空间了，至于 labels 文件夹下只能自己生成，这里本人提供一种从 VOC格式的 Annotations转Yolo5的labels代码：

3 YOLO 5 实验环境配置

YOLO5中 官网给的 install 步骤很简单，即直接 pip install -U -r requirements.txt ，但是这样会下载一些库的最新版本，如果您的实验环境中还有其他程序，像我只有一块GPU，也实在不想因为某些库的变化，导致其他程序跑不了了，所以，本人的步骤大致为下：
1 创建虚拟环境：
conda create -n yolo5 python=3.7
2 切换至当前的 yolo5环境
conda activate yolo5
3 如果您放心，可以直接使用 pip install -U requirements.txt
如果您不放心，那就将 requirements.txt中的库一条条的安装，比如：
pip install Cython .
…
由于我的很多库之前都装好了，所以我就没再安装其他的库。
到此假设您的环境配置没问题，下面就是调试代码部分了。

4 训练自己的数据集

YOLO5的代码写的其实挺不错，我们在训练自己的数据集过程中，只需要改变部分参数就行了，下面一一来讲。

4.1 yolov5-2.0\data文件夹

该文件夹目录如下，这里主要是对训练和测试数据的一个基本参数配置，包含了您数据集的 images文件夹，类的种数和具体类标 其实我们主要是新建一个yaml文件，假设命名为 mydata.yaml，格式就仿照该文件夹下任意一个yaml文件即可，推荐仿照 coco128.yaml。

# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
train: ../YourDataDir/images/train/  # 128 images
val: ../YourDataDir/images/val/  # 128 images

# number of classes
nc: K

# class names
names: ['zhonglei_1','zhonglei_2','zhonglei_3','zhonglei_4',...,'zhonglei_K']

4.2 yolov5-2.0\models文件夹

这个文件夹主要存放的是 和模型相关的配置文件，一般情况下，就检测来讲，如果不是很难的数据集，YOLO5s 就会有很好的表现，这主要原因是YOLO5在训练时候运用了数据增强技术，貌似是"马赛克"（mosaic）增强，YOLO4貌似也是这个。
但是注意我们需要调整的参数主要是 类别数量 和 anchor 参数，其余参数也可以调整，不过一般建议默认吧，anchors 参数为您图像缩放为 640*640 对应对的大小，一般用KMeans进行聚类选取

Anchor 聚类选取的代码如下：
这个代码是 utils/utils.py 中的 一个方法，这个方法主要是假设您自己配置的 anchor 最高metric 评价值如果小于 0.99 ，该方法就会自动调用，帮您重新生成Anchors 自己可以看看琢磨下。

def kmean_anchors(path='./data/coco128.yaml', n=9, img_size=640, thr=4.0, gen=1000, verbose=True):
    """ Creates kmeans-evolved anchors from training dataset

        Arguments:
            path: path to dataset *.yaml, or a loaded dataset
            n: number of anchors
            img_size: image size used for training
            thr: anchor-label wh ratio threshold hyperparameter hyp['anchor_t'] used for training, default=4.0
            gen: generations to evolve anchors using genetic algorithm

        Return:
            k: kmeans evolved anchors

        Usage:
            from utils.utils import *; _ = kmean_anchors()
    """
    thr = 1. / thr

    def metric(k, wh):  # compute metrics
        r = wh[:, None] / k[None]
        x = torch.min(r, 1. / r).min(2)[0]  # ratio metric
        # x = wh_iou(wh, torch.tensor(k))  # iou metric
        return x, x.max(1)[0]  # x, best_x

    def fitness(k):  # mutation fitness
        _, best = metric(torch.tensor(k, dtype=torch.float32), wh)
        return (best * (best > thr).float()).mean()  # fitness

    def print_results(k):
        k = k[np.argsort(k.prod(1))]  # sort small to large
        x, best = metric(k, wh0)
        bpr, aat = (best > thr).float().mean(), (x > thr).float().mean() * n  # best possible recall, anch > thr
        print('thr=%.2f: %.4f best possible recall, %.2f anchors past thr' % (thr, bpr, aat))
        print('n=%g, img_size=%s, metric_all=%.3f/%.3f-mean/best, past_thr=%.3f-mean: ' %
              (n, img_size, x.mean(), best.mean(), x[x > thr].mean()), end='')
        for i, x in enumerate(k):
            print('%i,%i' % (round(x[0]), round(x[1])), end=',  ' if i < len(k) - 1 else '\n')  # use in *.cfg
        return k

    if isinstance(path, str):  # *.yaml file
        with open(path) as f:
            data_dict = yaml.load(f, Loader=yaml.FullLoader)  # model dict
        from utils.datasets import LoadImagesAndLabels
        dataset = LoadImagesAndLabels(data_dict['train'], augment=True, rect=True)
    else:
        dataset = path  # dataset

    # Get label wh
    shapes = img_size * dataset.shapes / dataset.shapes.max(1, keepdims=True)
    wh0 = np.concatenate([l[:, 3:5] * s for s, l in zip(shapes, dataset.labels)])  # wh

    # Filter
    i = (wh0 < 3.0).any(1).sum()
    if i:
        print('WARNING: Extremely small objects found. '
              '%g of %g labels are < 3 pixels in width or height.' % (i, len(wh0)))
    wh = wh0[(wh0 >= 2.0).any(1)]  # filter > 2 pixels

    # Kmeans calculation
    from scipy.cluster.vq import kmeans
    print('Running kmeans for %g anchors on %g points...' % (n, len(wh)))
    s = wh.std(0)  # sigmas for whitening
    k, dist = kmeans(wh / s, n, iter=30)  # points, mean distance
    k *= s
    wh = torch.tensor(wh, dtype=torch.float32)  # filtered
    wh0 = torch.tensor(wh0, dtype=torch.float32)  # unflitered
    k = print_results(k)

    # Plot
    # k, d = [None] * 20, [None] * 20
    # for i in tqdm(range(1, 21)):
    #     k[i-1], d[i-1] = kmeans(wh / s, i)  # points, mean distance
    # fig, ax = plt.subplots(1, 2, figsize=(14, 7))
    # ax = ax.ravel()
    # ax[0].plot(np.arange(1, 21), np.array(d) ** 2, marker='.')
    # fig, ax = plt.subplots(1, 2, figsize=(14, 7))  # plot wh
    # ax[0].hist(wh[wh[:, 0]<100, 0],400)
    # ax[1].hist(wh[wh[:, 1]<100, 1],400)
    # fig.tight_layout()
    # fig.savefig('wh.png', dpi=200)

    # Evolve
    npr = np.random
    f, sh, mp, s = fitness(k), k.shape, 0.9, 0.1  # fitness, generations, mutation prob, sigma
    pbar = tqdm(range(gen), desc='Evolving anchors with Genetic Algorithm')  # progress bar
    for _ in pbar:
        v = np.ones(sh)
        while (v == 1).all():  # mutate until a change occurs (prevent duplicates)
            v = ((npr.random(sh) < mp) * npr.random() * npr.randn(*sh) * s + 1).clip(0.3, 3.0)
        kg = (k.copy() * v).clip(min=2.0)
        fg = fitness(kg)
        if fg > f:
            f, k = fg, kg.copy()
            pbar.desc = 'Evolving anchors with Genetic Algorithm: fitness = %.4f' % f
            if verbose:
                print_results(k)

    return print_results(k)

metric 方法如下：

def metric(k):  # compute metric
        r = wh[:, None] / k[None]
        x = torch.min(r, 1. / r).min(2)[0]  # ratio metric
        best = x.max(1)[0]  # best_x
        return (best > 1. / thr).float().mean()  #  best possible recall

    bpr = metric(m.anchor_grid.clone().cpu().view(-1, 2))
    print('Best Possible Recall (BPR) = %.4f' % bpr, end='')
    if bpr < 0.99:  # threshold to recompute
        print('. Attempting to generate improved anchors, please wait...' % bpr)
        na = m.anchor_grid.numel() // 2  # number of anchors
        new_anchors = kmean_anchors(dataset, n=na, img_size=imgsz, thr=thr, gen=1000, verbose=False)
        new_bpr = metric(new_anchors.reshape(-1, 2))

4.3 train.py 调整

注意看 trian.py 中的 198 行，将 80 调整为 您的种类数，这个是个分类的超参数，应该是分类损失的权重，不调整其实也可以，最好把他设置默认为 1

另外需要调整的就是我画横线的超参数，这些调整为自己配置的，就可以开始训练了。

4.4 训练模型

yolo5s.pt是从官网下载的预训练模型，如果不写 ./weights/yolo5s.pt 直接按照官网上的 --weigths yolo5s 这就会下载 yolo5s 模型，由于我们使用的是 2.0版本，直接下载可能下载不下来，所以自己最好提前下载好，放在 weights 文件夹下

python train.py --weights ./weights/yolo5s.pt

【提示】一般来讲，训练10轮以上 模型就会出现明显的收敛，如果没有出现，最好关注下自己的数据集 或者 换个模型。

【上述内容中 产生 labels 还有不完善的地方,明天补充下】