Widerface数据集 | 用YOLOV5训练widerface数据集

github代码：https://github.com/mary-0830/Yolov5_Widerface

一、下载yolov5代码

源代码链接：https://github.com/ultralytics/yolov5

训练自己数据集：https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data

注意： 若出现fatal: unable to access ‘https://github.comxxxxxxxxxxx’: Failed to connect to xxxxxxxxxxxxx这类错误
解决方法： 将命令行里的http改为git重新执行。 链接

二、下载widerface数据集

下载链接：http://shuoyang1213.me/WIDERFACE/

使用google下载比较快，也可以用腾讯下载，还要把标注文件给下载下来。

三、把widerface数据集转成yolo格式

参考链接：https://blog.csdn.net/mary_0830/article/details/116589279

四、修改配置文件

在data/中添加一个widerface.yaml的配置文件，如下图所示，nc改为1，name改为"face"，train和val改为之前转好的yolo格式。

五、训练数据

将下图参数设置一下，然后执行 python train.py 即可。

开始训练

六、踩坑记录

1. requirements.txt 这个最好一个一个的安装，一起安装可能会出问题。
2. 出现runtimeError: Couldn't load custom C++ ops. This can happen if your PyTorch and torchvision versions are incompatible, or if you had errors while compiling torchvision from source.问题，使用下面的方法三解决。找到对应的版本。原链接
   
   下面是笔者的环境版本：
   

七、训练结果

测试结果：



标签相关图：

PR曲线：

精度曲线：

八、模型量化

参考：https://blog.csdn.net/weixin_44936889/article/details/110732476

可以看到原始模型有600MB，而进行模型量化后的模型只有166MB，压缩了很多空间。

源码：

# coding:utf-8
# 模型量化
import os
import torch

import torch
import torch.nn as nn
from tqdm import tqdm


def autopad(k, p=None):  
    # Pad to 'same'
    if p is None:
        p = k // 2 if isinstance(k, int) else [x // 2 for x in k]  # auto-pad
    return p


class Conv(nn.Module):
    # Standard convolution
    # ch_in, ch_out, kernel, stride, padding, groups
    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):
        super(Conv, self).__init__()
        self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p),
                              groups=g, bias=False)
        self.bn = nn.BatchNorm2d(c2)
        self.act = nn.Hardswish() if act else nn.Identity()

    def forward(self, x):
        return self.act(self.bn(self.conv(x)))

    def fuseforward(self, x):
        return self.act(self.conv(x))


class Ensemble(nn.ModuleList):
    # Ensemble of models
    def __init__(self):
        super(Ensemble, self).__init__()

    def forward(self, x, augment=False):
        y = []
        for module in self:
            y.append(module(x, augment)[0])
        # y = torch.stack(y).max(0)[0]  # max ensemble
        # y = torch.cat(y, 1)  # nms ensemble
        y = torch.stack(y).mean(0)  # mean ensemble
        return y, None  # inference, train output


def attempt_load(weights, map_location=None):

    model = Ensemble()
    for w in weights if isinstance(weights, list) else [weights]:
        # load FP32 model
        model.append(torch.load(
            w, map_location=map_location)['model'].float().fuse().eval())

    # Compatibility updates
    for m in tqdm(model.modules()):
        if type(m) in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6]:
            m.inplace = True  # pytorch 1.7.0 compatibility
        elif type(m) is Conv:
            m._non_persistent_buffers_set = set()  # pytorch 1.6.0 compatibility

    if len(model) == 1:
        return model[-1]  # return model
    else:
        print('Ensemble created with %s\n' % weights)
        for k in ['names', 'stride']:
            setattr(model, k, getattr(model[-1], k))
        return model  # return ensemble


def select_device(device='', batch_size=None):
    # device = 'cpu' or '0' or '0,1,2,3'
    cpu_request = device.lower() == 'cpu'
    if device and not cpu_request:  # if device requested other than 'cpu'
        os.environ['CUDA_VISIBLE_DEVICES'] = device  # set environment variable
        assert torch.cuda.is_available(
        ), 'CUDA unavailable, invalid device %s requested' % device  # check availablity

    cuda = False if cpu_request else torch.cuda.is_available()
    if cuda:
        c = 1024 ** 2  # bytes to MB
        ng = torch.cuda.device_count()
        if ng > 1 and batch_size:  # check that batch_size is compatible with device_count
            assert batch_size % ng == 0, 'batch-size %g not multiple of GPU count %g' % (
                batch_size, ng)
        x = [torch.cuda.get_device_properties(i) for i in range(ng)]
        s = f'Using torch {torch.__version__} '
        for i in range(0, ng):
            if i == 1:
                s = ' ' * len(s)

    return torch.device('cuda:0' if cuda else 'cpu')


if __name__ == '__main__':

    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument('--in_weights', type=str,
                        default='last.pt', help='initial weights path')
    parser.add_argument('--out_weights', type=str,
                        default='slim_model.pt', help='output weights path')
    parser.add_argument('--device', type=str, default='0', help='device')
    opt = parser.parse_args()

    device = select_device(opt.device)
    model = attempt_load(opt.in_weights, map_location=device)
    model.to(device).eval()
    model.half()

    torch.save(model, opt.out_weights)
    print('done.')

    print('-[INFO] before: {} kb, after: {} kb'.format(
        os.path.getsize(opt.in_weights), os.path.getsize(opt.out_weights)))

在命令行执行：

python slim.py --in_weights best.pt --out_weights slim_model.pt --device 0

笔者的话： 不得不说，Yolov5这个代码维护越来越棒了，方便小白训练。一直在用20年8月的版本，那个版本还是比较简单的，也没有两张显卡训练，也没有其他的日志保存，缺少的信息还是很多的。快一年过去了，这个版本已经非常成熟了，也用到了各个检测任务上，非常棒！~