rtmpose-m_210e_ap10k-288x288_p4p5_sacr.py

# Copyright (c) OpenMMLab. All rights reserved.
"""RTMPose-M with P4/P5 semantic-detail fusion and SACR on AP-10K.

This configuration implements the AP-10K Ours-M variant reported in
the paper:

    - AP-10K Split 1
    - input size: 288 x 288
    - P4/P5 semantic-detail fusion
    - stride-16 fused features for RTMCCHead
    - Scale-Adaptive Crop Resolution (SACR)
    - 210 training epochs
    - Stage-2 training during the final 30 epochs
    - flip test enabled for accuracy evaluation

MMPose follows the (width, height) convention for ``input_size``.
"""

_base_ = [
    'mmpose::animal_2d_keypoint/rtmpose/ap10k/'
    'rtmpose-m_8xb64-210e_ap10k-256x256.py'
]

# Import the custom modules provided by this repository.
custom_imports = dict(
    imports=['rtmpose_spatial_enhancement'],
    allow_failed_imports=False,
)

default_scope = 'mmpose'


# ============================================================
# Runtime settings
# ============================================================

max_epochs = 210
stage2_num_epochs = 30

# The actual AP-10K experiment used lr=5e-4.
base_lr = 5e-4

train_cfg = dict(
    max_epochs=max_epochs,
    val_interval=10,
)

randomness = dict(seed=21)


# ============================================================
# Optimization
# ============================================================

# The paper experiment used a physical batch size of 64.
# No gradient accumulation was used for this AP-10K experiment.
optim_wrapper = dict(
    _delete_=True,
    type='OptimWrapper',

    optimizer=dict(
        type='AdamW',
        lr=base_lr,
        weight_decay=0.05,
    ),

    clip_grad=dict(
        max_norm=35,
        norm_type=2,
    ),

    paramwise_cfg=dict(
        norm_decay_mult=0,
        bias_decay_mult=0,
        bypass_duplicate=True,
    ),
)

param_scheduler = [
    dict(
        type='LinearLR',
        start_factor=1.0e-5,
        by_epoch=False,
        begin=0,
        end=1000,
    ),
    dict(
        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,
    ),
]

auto_scale_lr = dict(
    base_batch_size=512,
)


# ============================================================
# Input and codec settings
# ============================================================

# MMPose convention: (width, height)
input_size = (288, 288)
num_keypoints = 17

codec = dict(
    type='SimCCLabel',
    input_size=input_size,
    sigma=(5.66, 5.66),
    simcc_split_ratio=2.0,
    normalize=False,
    use_dark=False,
)


# ============================================================
# Model
# ============================================================

model = dict(
    # Output both P4 and P5 features from CSPNeXt-M.
    backbone=dict(
        out_indices=(3, 4),
    ),

    # Lightweight P4/P5 semantic-detail fusion.
    neck=dict(
        type='P4P5AddFusionNeck',
        in_channels=(384, 768),
        out_channels=384,
    ),

    # RTMCCHead receives the stride-16 fused feature.
    #
    # 288 / 16 = 18
    head=dict(
        in_channels=384,
        input_size=input_size,
        in_featuremap_size=(18, 18),
        decoder=codec,
    ),

    # The reported AP-10K accuracy experiment used flip test.
    test_cfg=dict(
        flip_test=True,
    ),
)


# ============================================================
# SACR settings
# ============================================================

sacr_cfg = dict(
    type='ScaleAdaptiveBBoxScale',
    base_padding=1.25,
    small_thr=96.0,
    medium_thr=160.0,
    small_padding=1.15,
    medium_padding=1.15,
    large_padding=1.25,
)


# ============================================================
# Dataset settings
# ============================================================

dataset_type = 'AP10KDataset'
data_mode = 'topdown'
data_root = 'data/ap10k/'

backend_args = dict(
    backend='local',
)


# ============================================================
# Training pipeline
# ============================================================

train_pipeline = [
    dict(
        type='LoadImage',
        backend_args=backend_args,
    ),

    dict(
        type='GetBBoxCenterScale',
        padding=1.25,
    ),

    # Scale-Adaptive Crop Resolution.
    dict(**sacr_cfg),

    dict(
        type='RandomFlip',
        direction='horizontal',
    ),

    dict(type='RandomHalfBody'),

    dict(
        type='RandomBBoxTransform',
        scale_factor=[0.6, 1.4],
        rotate_factor=80,
    ),

    dict(
        type='TopdownAffine',
        input_size=input_size,
    ),

    dict(type='mmdet.YOLOXHSVRandomAug'),

    dict(
        type='Albumentation',
        transforms=[
            dict(
                type='Blur',
                p=0.1,
            ),
            dict(
                type='MedianBlur',
                p=0.1,
            ),
            dict(
                type='CoarseDropout',
                max_holes=1,
                max_height=0.4,
                max_width=0.4,
                min_holes=1,
                min_height=0.2,
                min_width=0.2,
                p=1.0,
            ),
        ],
    ),

    dict(
        type='GenerateTarget',
        encoder=codec,
    ),

    dict(type='PackPoseInputs'),
]


# ============================================================
# Stage-2 training pipeline
# ============================================================

train_pipeline_stage2 = [
    dict(
        type='LoadImage',
        backend_args=backend_args,
    ),

    dict(
        type='GetBBoxCenterScale',
        padding=1.25,
    ),

    # SACR remains enabled during Stage-2 training.
    dict(**sacr_cfg),

    dict(
        type='RandomFlip',
        direction='horizontal',
    ),

    dict(type='RandomHalfBody'),

    dict(
        type='RandomBBoxTransform',
        shift_factor=0.0,
        scale_factor=[0.75, 1.25],
        rotate_factor=60,
    ),

    dict(
        type='TopdownAffine',
        input_size=input_size,
    ),

    dict(type='mmdet.YOLOXHSVRandomAug'),

    dict(
        type='Albumentation',
        transforms=[
            dict(
                type='Blur',
                p=0.1,
            ),
            dict(
                type='MedianBlur',
                p=0.1,
            ),
            dict(
                type='CoarseDropout',
                max_holes=1,
                max_height=0.4,
                max_width=0.4,
                min_holes=1,
                min_height=0.2,
                min_width=0.2,
                p=0.5,
            ),
        ],
    ),

    dict(
        type='GenerateTarget',
        encoder=codec,
    ),

    dict(type='PackPoseInputs'),
]


# ============================================================
# Validation and test pipeline
# ============================================================

val_pipeline = [
    dict(
        type='LoadImage',
        backend_args=backend_args,
    ),

    dict(
        type='GetBBoxCenterScale',
        padding=1.25,
    ),

    # Exactly one SACR transform is applied.
    dict(**sacr_cfg),

    dict(
        type='TopdownAffine',
        input_size=input_size,
    ),

    dict(type='PackPoseInputs'),
]


# ============================================================
# Data loaders
# ============================================================

train_dataloader = dict(
    _delete_=True,
    batch_size=64,
    num_workers=10,
    persistent_workers=True,

    sampler=dict(
        type='DefaultSampler',
        shuffle=True,
    ),

    dataset=dict(
        type=dataset_type,
        data_root=data_root,
        data_mode=data_mode,
        ann_file='annotations/ap10k-train-split1.json',
        data_prefix=dict(
            img='data/',
        ),
        pipeline=train_pipeline,
    ),
)


val_dataloader = dict(
    _delete_=True,
    batch_size=32,
    num_workers=10,
    persistent_workers=True,
    drop_last=False,

    sampler=dict(
        type='DefaultSampler',
        shuffle=False,
        round_up=False,
    ),

    dataset=dict(
        type=dataset_type,
        data_root=data_root,
        data_mode=data_mode,
        ann_file='annotations/ap10k-val-split1.json',
        data_prefix=dict(
            img='data/',
        ),
        test_mode=True,
        pipeline=val_pipeline,
    ),
)


test_dataloader = dict(
    _delete_=True,
    batch_size=32,
    num_workers=10,
    persistent_workers=True,
    drop_last=False,

    sampler=dict(
        type='DefaultSampler',
        shuffle=False,
        round_up=False,
    ),

    dataset=dict(
        type=dataset_type,
        data_root=data_root,
        data_mode=data_mode,
        ann_file='annotations/ap10k-test-split1.json',
        data_prefix=dict(
            img='data/',
        ),
        test_mode=True,
        pipeline=val_pipeline,
    ),
)


# ============================================================
# Hooks
# ============================================================

default_hooks = dict(
    checkpoint=dict(
        save_best='coco/AP',
        rule='greater',
        max_keep_ckpts=1,
    ),
)


custom_hooks = [
    dict(
        type='EMAHook',
        ema_type='ExpMomentumEMA',
        momentum=0.0002,
        update_buffers=True,
        priority=49,
    ),
    dict(
        type='mmdet.PipelineSwitchHook',
        switch_epoch=max_epochs - stage2_num_epochs,
        switch_pipeline=train_pipeline_stage2,
    ),
]


# ============================================================
# Evaluation
# ============================================================

val_evaluator = dict(
    type='CocoMetric',
    ann_file=(
        data_root
        + 'annotations/ap10k-val-split1.json'
    ),
)


test_evaluator = dict(
    type='CocoMetric',
    ann_file=(
        data_root
        + 'annotations/ap10k-test-split1.json'
    ),
)