Skip to main content

Finetuning Models

在 COCO 数据集上预训练的检测器可以作为其他数据集的良好预训练模型，例如 CityScapes 和 KITTI 数据集。本教程为用户提供指导，将 Model zoo 中提供的模型用于其他数据集以获得更好的性能。

在一个新的数据集上微调模型有两个步骤。

  • 按照 Customize Datasets 的方式添加对新数据集的支持。
  • 按照本教程中的讨论，修改配置。

以 Cityscapes 数据集的 finetuning 过程为例，用户需要修改配置中的五个部分。

继承 base config​

_base_ = [
    '../_base_/models/mask_rcnn_r50_fpn.py',
    '../_base_/datasets/cityscapes_instance.py', '../_base_/default_runtime.py'
]

为了减轻负担，减少编写整个配置的错误，MMDetection V2.0 支持从多个现有的配置中继承配置。为了微调一个 Mask RCNN 模型，新的配置需要继承 _base_/models/mask_rcnn_r50_fpn.py 来建立模型的基本结构。要使用 Cityscapes 数据集，新配置也可以简单地继承_base_/datasets/cityscapes_instance.py。对于训练时间表等运行时设置，新的配置需要继承 _base_/default_runtime.py。这些配置在 configs 目录下，用户也可以选择编写整个内容而不是使用继承。

修改 head​

然后，新的配置需要根据新数据集的类数来修改头部。通过只改变 roi_head 中的 num_classes，除了最后的预测头之外，预训练模型的权重大部分都可以重复使用。

model = dict(
    pretrained=None,
    roi_head=dict(
        bbox_head=dict(
            type='Shared2FCBBoxHead',
            in_channels=256,
            fc_out_channels=1024,
            roi_feat_size=7,
            num_classes=8,
            bbox_coder=dict(
                type='DeltaXYWHBBoxCoder',
                target_means=[0., 0., 0., 0.],
                target_stds=[0.1, 0.1, 0.2, 0.2]),
            reg_class_agnostic=False,
            loss_cls=dict(
                type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0),
            loss_bbox=dict(type='SmoothL1Loss', beta=1.0, loss_weight=1.0)),
        mask_head=dict(
            type='FCNMaskHead',
            num_convs=4,
            in_channels=256,
            conv_out_channels=256,
            num_classes=8,
            loss_mask=dict(
                type='CrossEntropyLoss', use_mask=True, loss_weight=1.0))))

修改 Dataset​

用户可能还需要准备数据集并编写关于数据集的配置。MMDetection V2.0 已经支持 VOC、WIDER FACE、COCO 和 Cityscapes 数据集。

修改训练 schedule​

微调超参数与默认计划不同。它通常需要较小的学习率和较少的训练周期

# optimizer
# lr is set for a batch size of 8
optimizer = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0001)
optimizer_config = dict(grad_clip=None)
# learning policy
lr_config = dict(
    policy='step',
    warmup='linear',
    warmup_iters=500,
    warmup_ratio=0.001,
    step=[7])
# the max_epochs and step in lr_config need specifically tuned for the customized dataset
runner = dict(max_epochs=8)
log_config = dict(interval=100)

使用预训练模型​

为了使用预训练的模型，新的配置文件中，在 load_from 中加入预训练模型的链接。用户可能需要在训练前下载模型的权重，以避免训练时的下载时间。

load_from = 'https://download.openmmlab.com/mmdetection/v2.0/mask_rcnn/mask_rcnn_r50_caffe_fpn_mstrain-poly_3x_coco/mask_rcnn_r50_caffe_fpn_mstrain-poly_3x_coco_bbox_mAP-0.408__segm_mAP-0.37_20200504_163245-42aa3d00.pth'  # noqa