rtmdet_config.py

# 当前路径位于 mmdetection/data, 配置将写到 mmdetection/data 路径下
 
config_balloon = """
_base_ = '../configs/rtmdet/rtmdet_tiny_8xb32-300e_coco.py'
data_root = './data/balloon/'
# 非常重要
metainfo = {
    # 类别名，注意 classes 需要是一个 tuple，因此即使是单类，
    # 你应该写成 `balloon,` 很多初学者经常会在这犯错220, 20, 60
    'classes': ('balloon',),
    'palette': [
        (0, 255, 255),
    ]
}
num_classes = 1
 
# 训练 40 epoch
max_epochs = 40
# 训练单卡 bs= 16
train_batch_size_per_gpu = 16
# 可以根据自己的电脑修改
train_num_workers = 4
 
# 验证集 batch size 为 1
val_batch_size_per_gpu = 1
val_num_workers = 2
 
# RTMDet 训练过程分成 2 个 stage，第二个 stage 会切换数据增强 pipeline
num_epochs_stage2 = 5
 
# batch 改变了，学习率也要跟着改变， 0.004 是 8卡x32 的学习率
base_lr = 16 * 0.004 / (32*8)
 
# 采用 COCO 预训练权重
load_from = 'https://download.openmmlab.com/mmdetection/v3.0/rtmdet/rtmdet_tiny_8xb32-300e_coco/rtmdet_tiny_8xb32-300e_coco_20220902_112414-78e30dcc.pth'  # noqa
 
model = dict(
    # 考虑到数据集太小，且训练时间很短，我们把 backbone 完全固定
    # 用户自己的数据集可能需要解冻 backbone
    backbone=dict(frozen_stages=4),
    # 不要忘记修改 num_classes
    bbox_head=dict(dict(num_classes=num_classes)))
 
# 数据集不同，dataset 输入参数也不一样
train_dataloader = dict(
    batch_size=train_batch_size_per_gpu,
    num_workers=train_num_workers,
    pin_memory=False,
    dataset=dict(
        data_root=data_root,
        metainfo=metainfo,
        ann_file='train/coco.json',
        data_prefix=dict(img='train/')))
val_dataloader = dict(
    batch_size=val_batch_size_per_gpu,
    num_workers=val_num_workers,
    dataset=dict(
        metainfo=metainfo,
        data_root=data_root,
        ann_file='val/coco.json',
        data_prefix=dict(img='val/')))
test_dataloader = val_dataloader
# 默认的学习率调度器是 warmup 1000，但是 balloon 数据集太小了，需要修改 为 30 iter
param_scheduler = [
    dict(
        type='LinearLR',
        start_factor=1.0e-5,
        by_epoch=False,
        begin=0,
        end=30),
    dict(
        type='CosineAnnealingLR',
        eta_min=base_lr * 0.05,
        begin=max_epochs // 2,  # max_epoch 也改变了
        end=max_epochs,
        T_max=max_epochs // 2,
        by_epoch=True,
        convert_to_iter_based=True),
]
optim_wrapper = dict(optimizer=dict(lr=base_lr))
# 第二 stage 切换 pipeline 的 epoch 时刻也改变了
_base_.custom_hooks[1].switch_epoch = max_epochs - num_epochs_stage2
val_evaluator = dict(ann_file=data_root + 'val/coco.json')
test_evaluator = val_evaluator
# 一些打印设置修改
default_hooks = dict(
    checkpoint=dict(interval=10, max_keep_ckpts=2, save_best='auto'),  # 同时保存最好性能权重
    logger=dict(type='LoggerHook', interval=5))
train_cfg = dict(max_epochs=max_epochs, val_interval=10)
"""
 
with open('./data/rtmdet_tiny_1xb16-40e_balloon.py', 'w') as f:
    f.write(config_balloon)