简单易用的分类+检测框架simpleAICV

所有代码已上传到本人github repository：https://github.com/zgcr/simpleAICV-pytorch-ImageNet-COCO-training

如果觉得有用，请点个star哟！

简介

又有好几个月没更新了，这次更新主要是对repository的代码优化，使之成为一个简单易用的分类+检测框架，整个框架代码编写主要依赖pytorch和opencv两个库。
在目标检测领域，其实已经有不少框架，比如mmdetection等。这类框架上手很简单，能够很容易地复现其中支持的各种算法，但是当你要对某种算法做一些网络结构、ground truth分配、loss计算以及decode方式等方面作一些深度定制修改时，就会变得非常麻烦。甚至，当我想了解某种算法实现时的某个基础结构或是ground truth分配具体如何实现中，需要在该框架中查找七八个.py文件，但仍然很难看到最底层实现是什么样的（函数套了太多层），和该算法论文中介绍的方法是否有区别。
因此，我编写了这个简单的分类+检测框架simpleAICV，写这个框架的目的就是获得最大的自由定制性，在这个框架中可以很容易的找到某算法的某个部分的详细实现（查找不超过三层函数），并且可以自由地修改。同样，对于训练和测试的过程，也可以很容易地进行修改。
这次更新主要增加内容：

1. 对分类和检测任务，增加公用的train.py和test.py文件，所有分类和检测模型都可以通过这两个文件执行训练和测试；
2. 如果对模型进行了训练和测试，在模型config.py文件所在工作目录分别自动生成train.info.log和test.info.log文件；
3. 利用torchvision0.9.0提供的torchvision.ops.deform_conv2d编写DCNv2，不再需要从这里安装：https://github.com/CharlesShang/DCNv2.git。自行编写的DCNv2在RTX20系列和RTX30系列显卡上均可兼容，不需要调整机器的CUDA driver版本；
4. 编写yolov4和yolov5的loss、decoder，初步训练了10个epoch，可以正常收敛，但由于yolo系列网络本身由于ground truth分配原则loss收敛较慢，且要在COCO数据集上从头开始训练需要大量的epoch数（500-600个epoch），这里并没有训练出最终结果。

公用的train.py和test.py文件

公用的train.py和test.py文件可以在tools目录内找到，其中分类和检测任务都有独立的公用train.py和test.py文件。当不需要深度修改train.py文件和test.py文件时，可以直接使用公用train.py和test.py文件。当你想做一些定制化修改时，可以将公用train.py和test.py文件复制到你的工作目录下，只需要保证.py文件中的BASE_DIR为本框架主目录路径即可，然后就可以自由地修改train.py和test.py文件了。

path.py文件

path.py文件记录了训练时使用的数据集和预训练模型的存放地址，当把本框架移动到其他机器上时要对path.py文件进行修改，使其能够正确找到数据集和预训练模型。

seed的设置

在训练时为了使模型多次训练的结果不产生较大的波动，需要对训练和测试时的种子进行固定，这主要依靠tools/utils.py文件中的set_seed函数和worker_seed_init_fn函数实现。set_seed中固定了一系列的种子，为了保证运算效率，cudnn.benchmark和cudnn.deterministic分别被设置为False和True。worker_seed_init_fn主要是对各节点worker的种子进行设置，worker的种子对各个worker中作的随机数据增强的随机性会有影响，比如各类random操作、mosaic数据增强操作等等。

def worker_seed_init_fn(worker_id, num_workers, local_rank, seed):
    # worker_seed_init_fn function will be called at the beginning of each epoch
    # for each epoch the same worker has same seed value,so we add the current time to the seed
    worker_seed = num_workers * local_rank + worker_id + seed + int(
        time.time())
    np.random.seed(worker_seed)
    random.seed(worker_seed)

由于dataloader在每个epoch开始前都会重新初始化，为了使同编号worker在每个epoch开始时的随机种子发生变化，增强数据的随机性，在worker_seed中加入了当前时间作为随机数。

计算模型flops和params

这里我使用了第三方包thop，编写函数compute_flops_and_params实现该功能。需要注意的是，目前包括thop在内的所有可用于计算模型flops的第三方包均会在每一层命名额外的buffer变量"total_ops"和"total_params"，因此这个函数使用时一定要注意放在保存模型操作之后，否则你保存的模型中每一层也会有这两个buffer变量。

def compute_flops_and_params(config, model):
    flops_input = torch.randn(1, 3, config.input_image_size,
                              config.input_image_size)

    model_on_cuda = next(model.parameters()).is_cuda
    if model_on_cuda:
        flops_input = flops_input.cuda()

    flops, params = profile(model, inputs=(flops_input, ), verbose=False)
    flops, params = clever_format([flops, params], '%.3f')

    return flops, params

optimizer和scheduler

这两部分靠函数build_optimizer和build_scheduler实现。根据以往经验，optimizer只支持最常用的带momentum的SGD和AdamW，其他的优化方式在实际业务流程中很少使用。scheduler只支持MultiStepLR和CosineLR，两种方式均可以指定是否要使用warm up(按epoch warm up)。如果要使用build_scheduler函数，请确保Pytorch版本>=1.5，否则会报错。

def build_optimizer(config, model):
    assert config.optimizer in ['SGD', 'AdamW'], 'Unsupported optimizer!'

    if config.optimizer == 'SGD':
        return torch.optim.SGD(model.parameters(),
                               lr=config.lr,
                               momentum=config.momentum,
                               weight_decay=config.weight_decay)
    elif config.optimizer == 'AdamW':
        return torch.optim.AdamW(model.parameters(),
                                 lr=config.lr,
                                 weight_decay=config.weight_decay)


def build_scheduler(config, optimizer):
    '''
    The value of config.warm_up_epochs is zero or an integer larger than 0
    '''
    assert config.scheduler in ['MultiStepLR',
                                'CosineLR'], 'Unsupported scheduler!'
    assert config.warm_up_epochs >= 0, 'Illegal warm_up_epochs!'
    if config.warm_up_epochs > 0:
        if config.scheduler == 'MultiStepLR':
            lr_func = lambda epoch: epoch / config.warm_up_epochs if epoch <= config.warm_up_epochs else config.gamma**len(
                [m for m in config.milestones if m <= epoch])
        elif config.scheduler == 'CosineLR':
            lr_func = lambda epoch: epoch / config.warm_up_epochs if epoch <= config.warm_up_epochs else 0.5 * (
                math.cos(
                    (epoch - config.warm_up_epochs) /
                    (config.epochs - config.warm_up_epochs) * math.pi) + 1)
    elif config.warm_up_epochs == 0:
        if config.scheduler == 'MultiStepLR':
            lr_func = lambda epoch: config.gamma**len(
                [m for m in config.milestones if m <= epoch])
        elif config.scheduler == 'CosineLR':
            lr_func = lambda epoch: 0.5 * (math.cos(
                (epoch - config.warm_up_epochs) /
                (config.epochs - config.warm_up_epochs) * math.pi) + 1)

    return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lr_func)

训练模式

build_training_mode函数用于建立训练模式。本函数支持多种训练模式：

单机单卡：
1. DataParallel
2. DataParallel+apex(混合精度训练）
3. DistributedDataParallel
4. DistributedDataParallel+syncbn(跨卡同步bn，使用后batch_size减半)
5. DistributedDataParallel+apex(混合精度训练）
6. DistributedDataParallel+apex(混合精度训练）+syncbn(跨卡同步bn，使用后batch_size减半)
单机多卡：
1. DataParallel
2. DataParallel+apex(混合精度训练）
3. DistributedDataParallel
4. DistributedDataParallel+syncbn(跨卡同步bn，使用后batch_size减半)
5. DistributedDataParallel+apex(混合精度训练）
6. DistributedDataParallel+apex(混合精度训练）+syncbn(跨卡同步bn，使用后batch_size减半)

由于本人没有多机多卡环境，无法编写测试多机多卡如何建立上述训练模式。上述训练模式中最常用的还是DataParallel+apex(混合精度训练)和DistributedDataParallel+apex(混合精度训练)，syncbn(跨卡同步bn)使用后会导致显卡占用加倍，但训练精度没有明显的优势，故一般不使用。build_training_mode函数实现如下：

def build_training_mode(config, model, optimizer):
    '''
    Choose model training mode:nn.DataParallel/nn.parallel.DistributedDataParallel,use apex or not
    '''
    if config.distributed:
        if config.sync_bn:
            model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).cuda()
        if config.apex:
            amp.register_float_function(torch, 'sigmoid')
            amp.register_float_function(torch, 'softmax')
            amp.register_float_function(torchvision.ops, 'deform_conv2d')
            model, optimizer = amp.initialize(model, optimizer, opt_level='O1')
            model = apex.parallel.DistributedDataParallel(model,
                                                          delay_allreduce=True)
            if config.sync_bn:
                model = apex.parallel.convert_syncbn_model(model).cuda()
        else:
            local_rank = torch.distributed.get_rank()
            model = nn.parallel.DistributedDataParallel(
                model, device_ids=[local_rank], output_device=local_rank)
    else:
        if config.apex:
            model, optimizer = amp.initialize(model, optimizer, opt_level='O1')

        model = nn.DataParallel(model)

    return model

注意切换使用DataParallel模式和DistributedDataParallel模式时除了修改train_config.py文件中相关项外，还要修改train.sh脚本。

train.info.log和test.info.log

这两个log文件由tools/utils.py文件中的get_logger函数生成。两个log会记录训练和测试时的各项超参数及结果，具体记录了什么请看公用的train.py和test.py文件中的logger.info()语句。get_logger函数实现如下：

def get_logger(name, log_dir):
    '''
    Args:
        name(str): name of logger
        log_dir(str): path of log
    '''

    if not os.path.exists(log_dir):
        os.makedirs(log_dir)

    logger = logging.getLogger(name)
    logger.setLevel(logging.INFO)

    info_name = os.path.join(log_dir, '{}.info.log'.format(name))
    info_handler = logging.handlers.TimedRotatingFileHandler(info_name,
                                                             when='D',
                                                             encoding='utf-8')
    info_handler.setLevel(logging.INFO)

    formatter = logging.Formatter('%(asctime)s - %(message)s',
                                  datefmt='%Y-%m-%d %H:%M:%S')

    info_handler.setFormatter(formatter)

    logger.addHandler(info_handler)

    return logger

利用torchvision0.9.0提供的torchvision.ops.deform_conv2d编写DCNv2

一般来说，大家使用DCNv2都是安装这个库：https://github.com/CharlesShang/DCNv2.git。但是这个库需要编译后才能使用，而编译时对机器安装的CUDA driver版本有要求，这点在使用了RTX30系列显卡上的机器上表现的非常明显，需要仔细地调整机器环境。为了避免编译问题，我使用了torchvision.ops.deform_conv2d来编写DCNv2。注意要使用我编写的DCNv2，torchvision版本需要>=0.9.0，否则会报错。根据 https://github.com/CharlesShang/DCNv2.git 库中对DCNv2的实现，我编写的DCNv2如下：

import math

import torch
import torch.nn as nn
import torchvision.ops

class DeformableConv2d(nn.Module):
    def __init__(self,
                 in_channels,
                 out_channels,
                 kernel_size=(3, 3),
                 stride=1,
                 padding=1,
                 groups=1,
                 bias=False):

        super(DeformableConv2d, self).__init__()
        self.padding = padding

        self.weight = nn.Parameter(
            torch.Tensor(out_channels, in_channels, kernel_size[0],
                         kernel_size[1]))
        self.bias = nn.Parameter(torch.Tensor(out_channels))

        self.offset_conv = nn.Conv2d(in_channels,
                                     2 * groups * kernel_size[0] *
                                     kernel_size[1],
                                     kernel_size=kernel_size,
                                     stride=stride,
                                     padding=self.padding,
                                     bias=True)

        self.mask_conv = nn.Conv2d(in_channels,
                                   1 * groups * kernel_size[0] *
                                   kernel_size[1],
                                   kernel_size=kernel_size,
                                   stride=stride,
                                   padding=self.padding,
                                   bias=True)

        n = in_channels
        for k in kernel_size:
            n *= k
        stdv = 1. / math.sqrt(n)
        self.weight.data.uniform_(-stdv, stdv)
        self.bias.data.zero_()

        nn.init.constant_(self.offset_conv.weight, 0.)
        nn.init.constant_(self.offset_conv.bias, 0.)

        nn.init.constant_(self.mask_conv.weight, 0.)
        nn.init.constant_(self.mask_conv.bias, 0.)

    def forward(self, x):
        offset = self.offset_conv(x)
        mask = torch.sigmoid(self.mask_conv(x))

        x = torchvision.ops.deform_conv2d(input=x,
                                          offset=offset,
                                          weight=self.weight,
                                          bias=self.bias,
                                          padding=self.padding,
                                          mask=mask)

        return x

注意torchvision.ops.deform_conv2d不能使用半精度计算，因此在建立训练模式时使用 amp.register_float_function(torchvision.ops, ‘deform_conv2d’)将该op强制规定使用单浮点精度计算。

编写yolov4和yolov5

先说说yolov4和yolov3的区别：训练trick这一块就不说了，请自行查看yolov4论文和官方代码；网络结构yolov4和yolov3区别很大，请查看yolov4论文和官方代码；在ground truth分配方面，yolov4和yolov3是完全一样的；loss计算这块，yolov4将box loss换成了CIoU loss，其他部分不变；decode这块，yolov4使用DIoU nms代替IoU NMS，其他部分不变。
yolov5和yolov4方面的区别：训练trick不完全一样；网络结构yolov5和yolov4也不一样；ground truth分配不一样；loss计算这块obj loss、reg loss、 cls loss用的计算方法是一样的，但各loss的具体超参数有区别；decode这块yolov5使用IoU NMS，其他和yolov4一样。
yolov4和yolov5具体实现代码可在simpleAICV/detection/目录内各.py文件中找到。对于各种训练的trick，我并没有使用，因为这些trick一般不具有泛化性，在COCO数据集上虽然涨点了，但换一个数据集可能就会掉点。目前使用的数据增强就是RandomHorizontalFlip、Normalize和YoloStyleResize（带multi scale），此外mosaic数据增强也可以使用，yolov5中对mosaic数据增强的实现就是直接随机取4张图全部resize然后拼起来，因此mosaic数据增强一般不会导致掉点。
由于本人计算资源非常有限，因此没有把yolo训练完，只训练了使用yolov4 loss和decode的yolov3（不使用任何COCO或ImageNet预训练参数），不使用mosaic数据增强，在resize=416时训练到200个epoch，mAP0.5:0.95达到0.121，且loss仍然在稳步下降，由于yolo系列的ground truth分配机制原因，估计至少要训练500到600个epoch才能得到最终结果。

为什么不用Pytorch自带的混合精度训练？

Pytorch1.6已经自带了混合精度训练，但是其autocast API兼容性不太好，在多线程情况下，autocast不能正确应用，需要在模型定义代码的forward函数前加上@autocast()装饰器才能解决这个问题，这样就比较麻烦。考虑到apex+pytorch1.8目前运行良好，并没有什么bug，所以还是继续使用apex。