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

用YOLOV5训练widerface数据集

  • 一、下载yolov5代码
  • 二、下载widerface数据集
  • 三、把widerface数据集转成yolo格式
  • 四、修改配置文件
  • 五、训练数据
  • 六、踩坑记录
  • 七、训练结果
  • 八、模型量化

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/
Widerface数据集 | 用YOLOV5训练widerface数据集_第1张图片
使用google下载比较快,也可以用腾讯下载,还要把标注文件给下载下来。

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

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

四、修改配置文件

在data/中添加一个widerface.yaml的配置文件,如下图所示,nc改为1,name改为"face",train和val改为之前转好的yolo格式。
Widerface数据集 | 用YOLOV5训练widerface数据集_第2张图片

五、训练数据

将下图参数设置一下,然后执行 python train.py 即可。
Widerface数据集 | 用YOLOV5训练widerface数据集_第3张图片
开始训练
Widerface数据集 | 用YOLOV5训练widerface数据集_第4张图片

六、踩坑记录

  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.问题,使用下面的方法三解决。找到对应的版本。原链接
    Widerface数据集 | 用YOLOV5训练widerface数据集_第5张图片
    下面是笔者的环境版本:
    Widerface数据集 | 用YOLOV5训练widerface数据集_第6张图片

七、训练结果

测试结果:
Widerface数据集 | 用YOLOV5训练widerface数据集_第7张图片
Widerface数据集 | 用YOLOV5训练widerface数据集_第8张图片
Widerface数据集 | 用YOLOV5训练widerface数据集_第9张图片
标签相关图:
Widerface数据集 | 用YOLOV5训练widerface数据集_第10张图片
PR曲线:
Widerface数据集 | 用YOLOV5训练widerface数据集_第11张图片
精度曲线:
Widerface数据集 | 用YOLOV5训练widerface数据集_第12张图片

八、模型量化

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

Widerface数据集 | 用YOLOV5训练widerface数据集_第13张图片
可以看到原始模型有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

Widerface数据集 | 用YOLOV5训练widerface数据集_第14张图片

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

你可能感兴趣的:(#,航拍,&,人脸数据集)