本人是从Tensorflow过来的,刚开始接触到MMDetection框架是有点懵的,因为这个框架在Pytorch基础上封装了好几层,这么做的好处是模块之间的耦合度很低,改动起来非常方便,但坏处是对于我这样的新手(对Pytorch也不甚了解),乍一看很难理解整个框架的运行流程,甚至都不知道如何查看对应的源码,更别说是从头搭建自己的网络了。


1. MMDetection是什么

  MMDetection是OpenMMLab家族中的一员,主要负责2D目标测领域(比如MMDetection3D则负责3D目标检测)。首先我们需要知道为什么会出现MMDetection这个框架。当前目标检测算法众多,方法复杂,细节较多,个人复现起来难度很大,而且由于缺少共享平台和统一规范,就算有人成功实现了某一个算法,也很难被其他人复用。

  于是商汤和港中文大学集中了一批人,使用统一的代码规范复现了当前大部分主流和前沿的模型,比如Faster R-CNN系列、YOLO系列,以及较新的DETR等(如下图所示),并提供了预训练模型。其他人只需要遵循这个规范,就能直接“白嫖”,不需要自己再重新实现一遍,而这个规范就是MMDetection。在丰富模型的基础上,MMDetection还支持自定义的扩展,可以在已有模型上进行修改,也可以自己从头搭建一个全新的模型,基本可以满足学术研究和工业落地的需求。

mmdetection test速度慢 mmdetection loadfrom_MMDetection

2. 整体算法流程

  所有的目标检测算法都可以按照训练和测试流程抽象成若干个模块,对于初学者来说只要理解各个模块的输入输出以及实现的功能即可,本篇博客不会进行展开,模块内部的实现逻辑之后会单独开博客分析。这个流程也对应框架的代码构建流程,所以理解这副图很重要。

mmdetection test速度慢 mmdetection loadfrom_1024程序员节_02

2.1 训练流程

  训练流程包含9个核心组件,但不是每个算法都需要,具体如下表所示。

模块名称

必需

功能

Backbone


【特征提取】比如ResNet系列

Neck


【特征增强】对Backbone的特征进行融合和增强,比如FPN

Head


【特征解码】目标检测网络最重要的部分,Head对特征图进行解码,得到算法期望的输出,比如目标框的类别和坐标,按照one-stage和two-stage可以分为DenseHead和RoIHead

BBox Assigner


【正负样本分配】由于检测网络输出的目标个数和真值目标大多数情况下是不匹配的,因此首先要进行正负样本分配,不同的正负样本分配策略会带来显著的性能差异,该模块至关重要

BBox Sampler


【正负样本平衡】一般目标检测中真值目标个数都是非常少的,正负样本比远小于1,为了避免数据极度不平衡引起的过拟合,需要适当对正负样本进行采样,平衡正负样本的数量

BBox Encoder


【编码变换】为了更好的收敛和平衡过个loss,将网络输出结果进行特定的编码变换,比如归一化,Encoder的输出可以认为是模型前向过程的最终输出

Loss


【Loss计算】检测网络一般分为分类loss和回归loss,提供模型迭代优化的依据

Enhance


【特征增强】一般指即插即用、能够对特征进行增强的模块,比如Dropout、Dropblock等

Training Tricks


【训练技巧】即我们熟知的模型调参方法,比如早停、学习率调整等

2.2 测试流程

  相较于训练流程,测试时只有模型的前向推理过程,因此不需要正负样本分配、平衡、计算loss等操作,流程会更简单一些。下表是测试流程特有的模块。

模块名称

必需

功能

BBox Decoder


【解码变换】对应测试流程中的BBox Encoder模块,训练时怎么对目标进行编码,测试时就怎么进行解码

BBox PostProcess


【后处理】得到目标框后,可能会出现重叠情况,所以一般需要根据IOU或置信度对输出目标进行处理,最常用到的是NMS方法

Training Tricks


【测试技巧】比如模型集成、多尺度测试等

3. 算法搭建流程

  以训练流程为例,对于Tensorflow和Pytorch,我们需要编写数据读取、数据预处理、数据增强、算法模型、loss函数、训练策略的代码,最后将其整合进train()函数中开始训练,过程十分繁琐。由于MMDetection已经实现了上述步骤中的绝大部分方法,我们只需要调用现成的函数即可,具体是在Config文件中配置好相应方法的参数,并将Config文件传给MMDetection自带的train()函数,然后框架就会解析Config文件,自动调用配置好的方法,完成训练流程。所以在MMDetection上搭建一个算法,我们要做的事情只有3件:准备数据集、编写Config文件、调用框架自带的train.py开始训练。

mmdetection test速度慢 mmdetection loadfrom_目标检测_03

  我们先从MMDetection自带的RetinaNet开始,在COCO数据集上打通训练和测试流程。

3.1 准备数据集

  MMDetection已经实现了COCO数据集的处理,我们这里就直接使用COCO 2014数据集。下载好的数据集目录结构如下图所示,annotations文件夹中以json文件格式存放了标注数据,其中目标框的标注信息在instances文件中。

mmdetection test速度慢 mmdetection loadfrom_MMDetection_04

3.2 编写Config文件

  RetinaNet的配置文件位于MMDetection源码的./configs/retinanet路径下,打开目录会发现里面有很多Config文件,文件命名规则遵循:

{model}_[model setting]_{backbone}_{neck}_[norm setting]_[misc]_[gpu x batch_per_gpu]_{schedule}_{dataset}

其中大括号表示必选,大括号表示可选。比如我们接下来要使用的配置文件retinanet_r50_fpn_1x_coco.py含义就是:模型名称是RetinaNet,主干是ResNet50,Neck是FPN,训练12个Epoch(1个x是12,2个x就是24),使用COCO数据集。更详细的字段说明可以在官方说明文档中查阅。

mmdetection test速度慢 mmdetection loadfrom_mmdetection test速度慢_05

  但当我们打开配置文件retinanet_r50_fpn_1x_coco.py时,发现里面只有几行代码:

_base_ = [
    '../_base_/models/retinanet_r50_fpn.py',
    '../_base_/datasets/coco_detection.py',
    '../_base_/schedules/schedule_1x.py', '../_base_/default_runtime.py'
]
# optimizer
optimizer = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0001)

  其实MMDetection中配置文件是通过继承 + 修改的方式完成用户自定义配置文件的。配置文件开头的_base_ = list()表示你需要继承的配置文件,然后通过重写的方式完成对应属性的修改。如果想要查看完整的配置文件信息,不需要依据_base_一级一级往上查找,可以通过官方给定的工具./tools/misc/print_config.py来打印配置文件:

python ./tools/misc/print_config.py ./configs/retinanet/retinanet_r50_fpn_1x_coco.py

  然后就可以看到retinanet_r50_fpn_1x_coco.py对应的完整的配置文件内容,第二节提到的各个模块都可以在配置文件中找到对应的定义。配置文件由一串字典dict和变量的定义组成,经由Config.fromfile(filepath)函数加载后会返回一个Config类型的变量(MMCV的一个数据结构),然后MMDetection框架就能根据这个Config调用相关的build_detector()方法构建对应的模块。

  具体地,build_detector()方法首先会根据字典中的type找到对应的类(Class),这个类的类名就是type字符串的值,且这个类一定是事先注册(Registry) 好的,MMDetection能够根据type值查询到具体的类,否则就会报错。比如在下面配置文件中,model的type值为RetinaNet,我们可以在./mmdet/models/detectors/retinanet.py中找到定义。

@DETECTORS.register_module()		# 表示这个类已经注册
class RetinaNet(SingleStageDetector):
    """Implementation of `RetinaNet <https://arxiv.org/abs/1708.02002>`_"""

    def __init__(self,
                 backbone,
                 neck,
                 bbox_head,
                 train_cfg=None,
                 test_cfg=None,
                 pretrained=None,
                 init_cfg=None):
        super(RetinaNet, self).__init__(backbone, neck, bbox_head, train_cfg,
                                        test_cfg, pretrained, init_cfg)

  我们可以发现RetinaNet类构造函数的参数刚好和配置文件中type='RetinaNet'的字典的其他键值对应。所以build_detector()函数的作用就是根据dict中的type找到对应的类,然后使用dict中传入的参数来对类进行初始化操作,并返回这个类的句柄。

# 下面两行调用是等价的
model = build_detector(Config{type='RetinaNet', backbone=xxx, neck=xxx, bbox_head=xxx})
model = RetinaNet(backbone=xxx, neck=xxx, bbox_head=xxx)

  然后配置文件中的dict是可以嵌套的,比如说model的backbone属性是type='ResNet'一个字典,同理我们也可以在./mmdet/models/backbones/resnet.py中找到ResNet类的定义,并且字典的键值和构造函数匹配。

@BACKBONES.register_module()
class ResNet(BaseModule):
    """ResNet backbone."""

    def __init__(self,
                 depth,
                 in_channels=3,
                 stem_channels=None,
                 base_channels=64,
                 num_stages=4,
                 strides=(1, 2, 2, 2),
                 dilations=(1, 1, 1, 1),
                 out_indices=(0, 1, 2, 3),
                 style='pytorch',
                 deep_stem=False,
                 avg_down=False,
                 frozen_stages=-1,
                 conv_cfg=None,
                 norm_cfg=dict(type='BN', requires_grad=True),
                 norm_eval=True,
                 dcn=None,
                 stage_with_dcn=(False, False, False, False),
                 plugins=None,
                 with_cp=False,
                 zero_init_residual=True,
                 pretrained=None,
                 init_cfg=None):
        super(ResNet, self).__init__(init_cfg)
        self.zero_init_residual = zero_init_residual
        if depth not in self.arch_settings:
            raise KeyError(f'invalid depth {depth} for resnet')

  下面是retinanet_r50_fpn_1x_coco.py完整的配置文件信息。

Config:
# 1. 模型配置
model = dict(
    type='RetinaNet',		# 模型名称
    # 1.1 Backbone配置
    backbone=dict(
        type='ResNet',		# Backbone使用ResNet50(4阶段,50层)
        depth=50,			
        num_stages=4,
        out_indices=(0, 1, 2, 3),	# 输出ResNet50第1~4阶段的feature map,供后续FPN做多尺度特征融合
        frozen_stages=1,			# 由于使用了预训练模型,冻结ResNet50第一阶段的网络参数,不参与训练过程
        norm_cfg=dict(type='BN', requires_grad=True),	# 归一化层配置
        norm_eval=True,
        style='pytorch',
        init_cfg=dict(type='Pretrained', checkpoint='torchvision://resnet50')),		# 使用pytorch提供的ResNet50在ImageNet上的预训练模型
    # 1.2 Neck配置
    neck=dict(
        type='FPN',		# Neck使用FPN
        in_channels=[256, 512, 1024, 2048],		# 输入通道数对应resnet50四个阶段feature map的维度
        out_channels=256,						# 输出特征维度为256
        start_level=1,							# 从Backbone的第一阶段特征图开始
        add_extra_convs='on_input',
        num_outs=5),
    # 1.3 Head配置
    bbox_head=dict(
        type='RetinaHead',	# Head使用RetinaHead
        num_classes=80,		# COCO数据集包含80类目标
        in_channels=256,	# FPN层输出特征维度为256
        stacked_convs=4,
        feat_channels=256,
        # 1.3.1 Retina是Anchor-Based方法, 需要生成Anchor
        anchor_generator=dict(
            type='AnchorGenerator',
            octave_base_scale=4,
            scales_per_octave=3,
            ratios=[0.5, 1.0, 2.0],
            strides=[8, 16, 32, 64, 128]),
        # 1.3.2 BBox Encoder配置
        bbox_coder=dict(
            type='DeltaXYWHBBoxCoder',
            target_means=[0.0, 0.0, 0.0, 0.0],
            target_stds=[1.0, 1.0, 1.0, 1.0]),
        # 1.3.3 分类Loss函数
        loss_cls=dict(
            type='FocalLoss',
            use_sigmoid=True,
            gamma=2.0,
            alpha=0.25,
            loss_weight=1.0),
        # 1.3.4 回归Loss函数
        loss_bbox=dict(type='L1Loss', loss_weight=1.0)),
    # 1.4 训练配置
    train_cfg=dict(
    	# 1.4.1 BBox Assigner
        assigner=dict(
            type='MaxIoUAssigner',
            pos_iou_thr=0.5,
            neg_iou_thr=0.4,
            min_pos_iou=0,
            ignore_iof_thr=-1),
        allowed_border=-1,
        pos_weight=-1,
        debug=False),
    # 1.5 测试配置
    test_cfg=dict(
        nms_pre=1000,
        min_bbox_size=0,
        score_thr=0.05,
        nms=dict(type='nms', iou_threshold=0.5),
        max_per_img=100))

# 2. 数据配置
data = dict(
    samples_per_gpu=2,		# batch_size大小
    workers_per_gpu=2,		# 每个GPU的线程数, 影响dataload的速度
    # 2.1 训练集配置
    train=dict(
        type='CocoDataset',
        ann_file='data/coco/annotations/instances_train2017.json',			
        img_prefix='data/coco/train2017/',									
        # 数据预处理步骤
        pipeline=[
            dict(type='LoadImageFromFile'),									
            dict(type='LoadAnnotations', with_bbox=True),					
            dict(type='Resize', img_scale=(1333, 800), keep_ratio=True),	
            dict(type='RandomFlip', flip_ratio=0.5),						
            dict(															
                type='Normalize',
                mean=[123.675, 116.28, 103.53],
                std=[58.395, 57.12, 57.375],
                to_rgb=True),
            dict(type='Pad', size_divisor=32),								
            dict(type='DefaultFormatBundle'),
            dict(type='Collect', keys=['img', 'gt_bboxes', 'gt_labels'])
        ]),
    # 2.2 验证集配置
    val=dict(
        type='CocoDataset',
        ann_file='data/coco/annotations/instances_val2017.json',
        img_prefix='data/coco/val2017/',
        pipeline=[
            dict(type='LoadImageFromFile'),
            dict(
                type='MultiScaleFlipAug',
                img_scale=(1333, 800),
                flip=False,
                transforms=[
                    dict(type='Resize', keep_ratio=True),
                    dict(type='RandomFlip'),
                    dict(
                        type='Normalize',
                        mean=[123.675, 116.28, 103.53],
                        std=[58.395, 57.12, 57.375],
                        to_rgb=True),
                    dict(type='Pad', size_divisor=32),
                    dict(type='ImageToTensor', keys=['img']),
                    dict(type='Collect', keys=['img'])
                ])
        ]),
    # 2.3 测试集配置
    test=dict(
        type='CocoDataset',
        ann_file='data/coco/annotations/instances_val2017.json',
        img_prefix='data/coco/val2017/',
        pipeline=[
            dict(type='LoadImageFromFile'),
            dict(
                type='MultiScaleFlipAug',
                img_scale=(1333, 800),
                flip=False,
                transforms=[
                    dict(type='Resize', keep_ratio=True),
                    dict(type='RandomFlip'),
                    dict(
                        type='Normalize',
                        mean=[123.675, 116.28, 103.53],
                        std=[58.395, 57.12, 57.375],
                        to_rgb=True),
                    dict(type='Pad', size_divisor=32),
                    dict(type='ImageToTensor', keys=['img']),
                    dict(type='Collect', keys=['img'])
                ])
        ]))
# evalution hook配置
evaluation = dict(interval=1, metric='bbox')
# 优化器配置
optimizer = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0001)
# optimizer hook配置
optimizer_config = dict(grad_clip=None)
# 学习率配置
lr_config = dict(
    policy='step',
    warmup='linear',
    warmup_iters=500,
    warmup_ratio=0.001,
    step=[8, 11])
# Runner配置
runner = dict(type='EpochBasedRunner', max_epochs=12)
# checkpoint配置
checkpoint_config = dict(interval=1)
# logger hook配置
log_config = dict(interval=50, hooks=[dict(type='TextLoggerHook')])
# 自定义hook配置
custom_hooks = [dict(type='NumClassCheckHook')]
# 分布式训练配置
dist_params = dict(backend='nccl')
# 日志级别
log_level = 'INFO'
# 预训练模型路径
load_from = None
# 模型断点
resume_from = None
# Runner的工作流
workflow = [('train', 1)]

  从配置文件可以看到,当前默认从pytorch官网下载预训练模型,且数据集的路径以及GPU数目和我当前的不符,而且由于电脑内存有限,我不希望每个epoch都保存一次checkpoint,所以我新建了一个配置文件my_retinanet_r50_fpn.py继承了官方的配置文件,并进行了一些修改:

_base_ = [
    'D:/Program Files/OpenSourceLib/mmdetection/configs/retinanet/retinanet_r50_fpn_1x_coco.py'
]

model = dict(
    backbone=dict(
        init_cfg=None)		# 不再直接从官网下载预训练模型,使用我自己下载好的预训练模型
)

data = dict(
    samples_per_gpu=2,		# batch_size=2
    workers_per_gpu=1,		# 每个GPU的线程数, 影响dataload的速度
    train=dict(
        type='CocoDataset',
        ann_file='E:/Dataset/COCO2014/annotations/instances_train2014.json',	# 修改数据集路径
        img_prefix='E:/Dataset/COCO2014/train2014'),
    val=dict(
        type='CocoDataset',
        ann_file='E:/Dataset/COCO2014/annotations/instances_val2014.json',
        img_prefix='E:/Dataset/COCO2014/val2014/'),
    test=dict(
        type='CocoDataset',
        ann_file='E:/Dataset/COCO2014/annotations/instances_val2014.json',
        img_prefix='E:/Dataset/COCO2014/val2014/')
)

evaluation = dict(interval=12, metric='bbox')	# 12个epoch进行一次评估

checkpoint_config = dict(interval=2)	# 2个epoch保存一次checkpoint

load_from = '../ckpts/resnet50-0676ba61.pth'	# 自己下载的预训练模型路径

3.3 训练网络

  写完配置文件后,就可以直接调用./tools/train.py指定配置文件进行训练。train.py包含了模型配置、数据集配置、训练配置、Hook配置等的解析,以及根据配置信息构造训练,用户的自定义操作可以通过Hook进行配置,一般无需修改train.py文件。

python train.py my_retinanet_r50_fpn.py

  成功开始训练:

mmdetection test速度慢 mmdetection loadfrom_深度学习_06

4. 总结

  本文利用MMDetection已经实现的RetinaNet模型在COCO上进行训练作为示例,演示了MMDetection的模型训练流程。总的来说分为三个步骤:

  1. 准备数据集
  2. 准备配置文件:配置文件由一系列dict组成,dict中的type键值代表注册的类别,build函数可以通过识别dict中的type来初始化对应的类。配置文件一般会继承一个通用配置文件,然后在此基础上根据需求调整。
  3. 开始训练:调用MMDetection自带的train.py进行训练。

  如果需要构建自己的模型,则需要实现一个类然后进行注册,Registry和Hook的机制见下一篇博客。