chenjie04's blog

RTMDet 是基于YOLOX再次开发的检测器，检测速度更快，精度更高，还可以做实例分割和旋转目标检测。

_base_ = [ # 继承默认配置
    '../_base_/default_runtime.py', '../_base_/schedules/schedule_1x.py',
    '../_base_/datasets/coco_detection.py', './rtmdet_tta.py'
]
model = dict(
    type='RTMDet',
    data_preprocessor=dict( # 数据预处理器相关配置
        # 数据预处理器的主要功能是把数据和模型从 CPU 搬运到 GPU，然后在GPU中对数据进行归一化操作，也可以实现MixUp、Mosaic 一类“将多张图片融合成一张图片”的数据增强操作，不过在这里设置batch_augments=None，即不做批量增强，因为RTMDet实现了新的CachedMosaic和CachedMixUp，采用缓存技术实现Mosaic和MixUp等需要多张图片的数据增广技术，把图片放在缓存中，训练时就不需要加载多张图片就可以实现Mosaic和MixUp，极大地提高数据加载速度。
        # 另外注意，根据mmdetection的设计原则，一般单张图像的数据增广流程会把它们组成train_pipeline传入到CocoDataset数据集类中，而多张图像的增广流程则在数据预处理器中实现，因为这里是DataLoader之后、模型之前，这里有一个batch的数据，可以很容易实现“将多张图片融合成一张图片”的数据增强
        type='DetDataPreprocessor',
        mean=[103.53, 116.28, 123.675], # 用于数据归一化的均值
        std=[57.375, 57.12, 58.395], # 用于数据归一化的标准差
        bgr_to_rgb=False, # 是否交换图片通道，这里为否
        batch_augments=None),
    backbone=dict( 
        # 基于5*5的深度可分离卷积重新设计了主干网络，采用深度可分离卷积必须额外接一个逐点卷积层 （point-wise convolution），这样模型的深度增加，因此设计团队适当减少主干网络的每个stage的基础模块并稍微增加模型宽度。
        type='CSPNeXt',
        arch='P5',
        expand_ratio=0.5,
        deepen_factor=1,
        widen_factor=1,
        channel_attention=True,
        norm_cfg=dict(type='SyncBN'),
        act_cfg=dict(type='SiLU', inplace=True)),
    neck=dict(
        # 一个更heavier的颈部网络用于融合多尺度特征对检测不同尺度的目标十分有必要。设计团队提高颈部网络基础模块的膨胀比，以将更多参数和计算放到颈部，提高模型能力。
        type='CSPNeXtPAFPN',
        in_channels=[256, 512, 1024],
        out_channels=256,
        num_csp_blocks=3,
        expand_ratio=0.5,
        norm_cfg=dict(type='SyncBN'),
        act_cfg=dict(type='SiLU', inplace=True)),
    bbox_head=dict(
        # 不同尺度特征图共享检测头，但是得配备不同的Batch Normalization (BN)层。
        type='RTMDetSepBNHead',
        num_classes=80,
        in_channels=256,
        stacked_convs=2,
        feat_channels=256,
        anchor_generator=dict(
            type='MlvlPointGenerator', offset=0, strides=[8, 16, 32]),
        bbox_coder=dict(type='DistancePointBBoxCoder'),
        loss_cls=dict(
            type='QualityFocalLoss',
            use_sigmoid=True,
            beta=2.0,
            loss_weight=1.0),
        loss_bbox=dict(type='GIoULoss', loss_weight=2.0),
        with_objectness=False,
        exp_on_reg=True,
        share_conv=True,
        pred_kernel_size=1,
        norm_cfg=dict(type='SyncBN'),
        act_cfg=dict(type='SiLU', inplace=True)),
    train_cfg=dict(
        # 重新设计了动态软标签分配策略
        assigner=dict(type='DynamicSoftLabelAssigner', topk=13),
        allowed_border=-1,
        pos_weight=-1,
        debug=False),
    test_cfg=dict( # 后处理
        nms_pre=30000, # NMS 前的 box 数
        min_bbox_size=0, # box 允许的最小尺寸
        score_thr=0.001, # bbox 的分数阈值
        nms=dict(type='nms', iou_threshold=0.65), # NMS 类型及阈值
        max_per_img=300), # NMS 后要保留的 box 数量
)

# 替换默认的train_pipeline，RTMDet采用两段式训练策略，通过PipelineSwitchHook实现
# 两段式训练是指将训练过程分成两个阶段，第一阶段采用强数据增广Mosaic和MixUp，第二阶段采用弱数据增广，即不用Mosaic和MixUp。YOLOX还在第二阶段引入L1损失函数来微调回归分支。在这里，设计团队不再引入L1损失函数，将第一阶段数据增广时混合图片数量提高到8，且第一阶段设置为280epochs；在最后20个epochs中采用 Large Scale Jittering (LSJ)做为数据增广。两段式训练的原因是强数据增广虽然能够提高模型的泛化性，但是会导致训练数据分布与真实分布不一致，所以需要在第二阶段使用与真实分布更接近的数据进行微调。
train_pipeline = [
    dict(
        type='LoadImageFromFile',
        file_client_args={{_base_.file_client_args}}),
    dict(type='LoadAnnotations', with_bbox=True),
    dict(type='CachedMosaic', img_scale=(640, 640), pad_val=114.0),
    dict(
        type='RandomResize',
        scale=(1280, 1280),
        ratio_range=(0.1, 2.0),
        keep_ratio=True),
    dict(type='RandomCrop', crop_size=(640, 640)),
    dict(type='YOLOXHSVRandomAug'),
    dict(type='RandomFlip', prob=0.5),
    dict(type='Pad', size=(640, 640), pad_val=dict(img=(114, 114, 114))),
    dict(
        type='CachedMixUp',
        img_scale=(640, 640),
        ratio_range=(1.0, 1.0),
        max_cached_images=20,
        pad_val=(114, 114, 114)),
    dict(type='PackDetInputs')
]

train_pipeline_stage2 = [
    dict(
        type='LoadImageFromFile',
        file_client_args={{_base_.file_client_args}}),
    dict(type='LoadAnnotations', with_bbox=True),
    dict(
        type='RandomResize',
        scale=(640, 640),
        ratio_range=(0.1, 2.0),
        keep_ratio=True),
    dict(type='RandomCrop', crop_size=(640, 640)),
    dict(type='YOLOXHSVRandomAug'),
    dict(type='RandomFlip', prob=0.5),
    dict(type='Pad', size=(640, 640), pad_val=dict(img=(114, 114, 114))),
    dict(type='PackDetInputs')
]

# 替换test_pipeline
test_pipeline = [
    dict(
        type='LoadImageFromFile',
        file_client_args={{_base_.file_client_args}}),
    dict(type='Resize', scale=(640, 640), keep_ratio=True),
    dict(type='Pad', size=(640, 640), pad_val=dict(img=(114, 114, 114))),
    dict(
        type='PackDetInputs',
        meta_keys=('img_id', 'img_path', 'ori_shape', 'img_shape',
                   'scale_factor'))
]

train_dataloader = dict(
    batch_size=32,
    num_workers=10,
    batch_sampler=None,
    pin_memory=True,
    dataset=dict(pipeline=train_pipeline))
val_dataloader = dict(
    batch_size=5, num_workers=10, dataset=dict(pipeline=test_pipeline))
test_dataloader = val_dataloader

max_epochs = 300
stage2_num_epochs = 20
base_lr = 0.004
interval = 10

train_cfg = dict(
    max_epochs=max_epochs,
    val_interval=interval,
    dynamic_intervals=[(max_epochs - stage2_num_epochs, 1)])

val_evaluator = dict(proposal_nums=(100, 1, 10))
test_evaluator = val_evaluator

# optimizer
optim_wrapper = dict(
    _delete_=True,
    type='OptimWrapper',
    optimizer=dict(type='AdamW', lr=base_lr, weight_decay=0.05),
    paramwise_cfg=dict(
        norm_decay_mult=0, bias_decay_mult=0, bypass_duplicate=True))

# learning rate
param_scheduler = [
    dict(
        type='LinearLR',
        start_factor=1.0e-5,
        by_epoch=False,
        begin=0,
        end=1000),
    dict(
        # use cosine lr from 150 to 300 epoch
        type='CosineAnnealingLR',
        eta_min=base_lr * 0.05,
        begin=max_epochs // 2,
        end=max_epochs,
        T_max=max_epochs // 2,
        by_epoch=True,
        convert_to_iter_based=True),
]

# hooks
default_hooks = dict(
    checkpoint=dict(
        interval=interval,
        max_keep_ckpts=3  # only keep latest 3 checkpoints
    ))
custom_hooks = [
    dict(
        type='EMAHook',
        ema_type='ExpMomentumEMA',
        momentum=0.0002,
        update_buffers=True,
        priority=49),
    dict(
        type='PipelineSwitchHook',
        switch_epoch=max_epochs - stage2_num_epochs,
        switch_pipeline=train_pipeline_stage2)
]