是时候该学会 MMDetection 进阶之非典型操作技能（一）

目录

前言

1 如何给不同 layer 设置不同的学习率以及冻结特定层

1.1 DefaultOptimizerConstructor

1.2 冻结特定层解决办法

2 如何在训练中优雅地使用多张图数据增强

3 如何在训练中实时调整数据预处理流程以及切换 loss

3.1 Dataloader 在开启多进程下无法实时修改内部属性解决办法

4 总结

---

前言

大家好，今天我们将开启全新的 MMDetection 系列文章，是时候带大家学习一些非典型操作技能啦。

这些非典型操作出现的原因各种各样，有部分来自内部和社区用户所提需求，有部分来自复现算法本身的需求。希望通过学习本系列文章，用户在使用 MMDetection 进行扩展开发时可以更加游刃有余，轻松秀出各种骚操作。

本文是非典型操作系列文章的首篇，所涉及到的典型操作技能为：

• 如何给不同 layer 设置不同的学习率以及冻结特定层
• 如何在训练中优雅地使用多图数据增强
• 如何在训练中实时调整数据预处理流程以及切换 loss

注意：
本文需要用户对 MMDetection 本身有一定了解，可通过  官方文档 或者  知乎文案 了解 MMDetection。
本文所述非典型操作办法可能仅适用于 MMDetection V2.21.0 及其以前版本 ，随着 MMDetection 持续更新，相信之后会有更优雅的解决办法。

1 如何给不同 layer 设置不同的学习率以及冻结特定层

经常看到 issue 中有人提到这个问题，其实 MMDetection 是支持给不同 layer 设置不同的学习率以及冻结特定层的，核心都是通过优化器构造器 Optimizer Constructor 实现的，MMCV 中提供了默认的

DefaultOptimizerConstructor 来处理用户平时能够遇到的大部分需求。

要给不同的层设置不同的学习率，可以参考 DETR 算法的 configs/detr/detr_r50_8x2_150e_coco.py 配置文件。

optimizer = dict( 
    type='AdamW', 
    lr=0.0001, 
    weight_decay=0.0001, 
    paramwise_cfg=dict( 
 custom_keys={'backbone': dict(lr_mult=0.1, decay_mult=1.0)})) 
 

上述配置的意思是给 DETR 算法中的 backbone 部分的初始化学习率全部乘上 0.1，也就是 backbone 学习率比 head 部分学习率小 10 倍。同样的，可以參考 Swin Transformer 算法的 configs/swin/mask_rcnn_swin-t-p4-w7_fpn_1x_coco.py 配置文件。

optimizer = dict( 
    _delete_=True, 
    type='AdamW', 
    lr=0.0001, 
    betas=(0.9, 0.999), 
    weight_decay=0.05, 
    paramwise_cfg=dict( 
 custom_keys={ 
            'absolute_pos_embed': dict(decay_mult=0.), 
            'relative_position_bias_table': dict(decay_mult=0.), 
            'norm': dict(decay_mult=0.) 
        })) 
 

将包含指定 key 的层的 decay 系数设为 0，也就是不进行 weight decay。

至于冻结特定层，目前只能用于无 BN 层的模块。幸好，大部分 FPN 和 Head 模块都是没有 BN 层的，所以大部分情况下用户都可以将想冻结的层中的 lr_mult 设置为0，从而间接达到目标。

1.1 DefaultOptimizerConstructor

首先要强调 OptimizerConstructor 的作用就是给不同层设置不同的模型优化超参，一般大家常见的配置是：

optimizer = dict(type='SGD', lr=0.02, momentum=0.9, weight_decay=0.0001) 

这表示所有层超参一视同仁，实际上构建优化器时候代码如下：

def build_optimizer(model, cfg): 
    optimizer_cfg = copy.deepcopy(cfg) 
    # 如果用户没有自定义 constructor ，则使用 DefaultOptimizerConstructor 
    constructor_type = optimizer_cfg.pop('constructor', 
                                         'DefaultOptimizerConstructor') 
    # 并取出 paramwise_cfg                                   
    paramwise_cfg = optimizer_cfg.pop('paramwise_cfg', None) 
    # 实例化 DefaultOptimizerConstructor 
    optim_constructor = build_optimizer_constructor( 
        dict( 
            type=constructor_type, 
            optimizer_cfg=optimizer_cfg, 
            paramwise_cfg=paramwise_cfg)) 
    # 返回 pytorch 的优化器对象        
    optimizer = optim_constructor(model) 
    return optimizer 

而 DefaultOptimizerConstructor 的示例代码为：

@OPTIMIZER_BUILDERS.register_module() 
class DefaultOptimizerConstructor: 
 
 def __init__(self, optimizer_cfg, paramwise_cfg=None): 
 
        self.optimizer_cfg = optimizer_cfg 
        self.paramwise_cfg = {} if paramwise_cfg is None else paramwise_cfg 
        # 这是优化器本身配置 
        self.base_lr = optimizer_cfg.get('lr', None) 
        self.base_wd = optimizer_cfg.get('weight_decay', None) 
 
    def add_params(self, params, module, prefix='', is_dcn_module=None): 
        # 这些参数很重要 
 bias_lr_mult = self.paramwise_cfg.get('bias_lr_mult', 1.) 
        bias_decay_mult = self.paramwise_cfg.get('bias_decay_mult', 1.) 
        norm_decay_mult = self.paramwise_cfg.get('norm_decay_mult', 1.) 
        dwconv_decay_mult = self.paramwise_cfg.get('dwconv_decay_mult', 1.) 
        bypass_duplicate = self.paramwise_cfg.get('bypass_duplicate', False) 
        dcn_offset_lr_mult = self.paramwise_cfg.get('dcn_offset_lr_mult', 1.) 
 
        for name, param in module.named_parameters(recurse=False): 
            param_group = {'params': [param]} 
            if not param.requires_grad: 
                params.append(param_group) 
                continue 
 
            # 对自定义 key 进行设置新的参数组参数 
            ... 
            # 添加到参数组 
            params.append(param_group) 
 
        # 遍历所有模块 
        for child_name, child_mod in module.named_children(): 
            child_prefix = f'{prefix}.{child_name}' if prefix else child_name 
            self.add_params( 
                params, 
                child_mod, 
                prefix=child_prefix, 
                is_dcn_module=is_dcn_module) 
 
    # 调用时候，返回 pytorch 优化器对象 
    def __call__(self, model): 
        optimizer_cfg = self.optimizer_cfg.copy() 
        # 如果 paramwise_cfg 參數沒有指定，則使用全局配置 
        if not self.paramwise_cfg: 
            optimizer_cfg['params'] = model.parameters() 
            return build_from_cfg(optimizer_cfg, OPTIMIZERS) 
 
        # 设置参数组 
        params = [] 
        self.add_params(params, model) 
        optimizer_cfg['params'] = params 
 
        return build_from_cfg(optimizer_cfg, OPTIMIZERS) 
 

从上面的参数可以知道，DefaultOptimizerConstructor 具备的功能为：

• bias_lr_mult 给特定层或者所有层的 bias_lr 乘上一个系数
• bias_decay_mult 给特定层或者所有层的 bias 模块的 decay 乘上一个系数
• 其他也是类似

当用户指定 custom_keys 时候，DefaultOptimizerConstructor 会遍历模型参数，然后通过字符串匹配方式查看 custom_keys 是否在模型参数中，如果在则会给当前参数组设置用户指定的系数。因为他是通过字符串匹配的方式判断，所以用户指定 custom_keys 时候要注意 key 的唯一性，否则可能出现额外匹配。例如用户只想给模型模块层 a.b.c 进行定制 lr，如果模型层还有名称为 a.b.d 的模块，此时用户设置 custom_key 为 a.b，那么就会同时匹配搭配 a.b.d了，此时就出现了额外匹配。简要核心代码实现如下：

 # 先按照字母表排序，然后按照长度反向排序，越短的在前 
 sorted_keys = sorted(sorted(custom_keys.keys()), key=len, reverse=True) 
 for name, param in module.named_parameters(recurse=False): 
     for key in sorted_keys: 
         if key in f'{prefix}.{name}': 
            lr_mult = custom_keys[key].get('lr_mult', 1.) 
            param_group['lr'] = self.base_lr * lr_mult 
            if self.base_wd is not None: 
                 decay_mult = custom_keys[key].get('decay_mult', 1.) 
                 param_group['weight_decay'] = self.base_wd * decay_mult 
                 break 

1.2 冻结特定层解决办法

对于没有 BN 层的模块，用户可以将想冻结的层中的 lr_mult 设置为 0。但是一旦有 BN，lr=0 只是可学习参数不再更新，但是全局均值和方差依然在改，没有实现真正的冻结。

如果面对有 BN 层的冻结需求时，暂时没有办法通过直接修改配置实现，目前来看有两种自定义办法：

• 自定义 OptimizerConstructor 或者继承 DefaultOptimizerConstructor，然后在内部自己处理逻辑
• 直接在构建模型的时候将想要冻结的层设置 requires_grad 属性为 False，并且切换为 eval 模式

对于一般水平用户推荐第二种最简单直接的做法，如果是有能力自定义 OptimizerConstructor 的用户则推荐直接自定义，这样写更加通用。

2 如何在训练中优雅地使用多张图数据增强

同时使用多张图数据增强的典型代表是 Mosaic 和 Mixup。Mosaic 数据增强一次会读 4 张图，每张图都要输入到训练增强 pipeline 中，并最终合并成 1 张大图输出。在支持 Mosaic 前，MMDetection 的 pipeline 不支持这种非典型范式，用户要想直接支持也比较困难。

基于扩展开发原则，我们希望在不大幅改动 MMDetection pipeline 的前提下能够支持多图数据增强，为此我们和 ConcatDataset 做法一样，新建了多图的 MultiImageMixDataset，代码位于 mmdet/datasets/dataset_wrappers.py。其核心实现为:

@DATASETS.register_module() 
class MultiImageMixDataset: 
    def __getitem__(self, idx): 
        results = copy.deepcopy(self.dataset[idx]) 
        for (transform, transform_type) in zip(self.pipeline, 
                                               self.pipeline_types): 
            # 如果当前 transform 中有 get_indexes 方法，则调用 
            if hasattr(transform, 'get_indexes'): 
                # 返回多张图片索引 
                indexes = transform.get_indexes(self.dataset) 
                if not isinstance(indexes, collections.abc.Sequence): 
                    indexes = [indexes] 
                # 然后获取多张图对应的原始数据 
                mix_results = [ 
                    copy.deepcopy(self.dataset[index]) for index in indexes 
                ] 
                results['mix_results'] = mix_results 
            # 再经过 transform，这样 transform 就可以一次性接收多张图片数据，从而同时进行增强和返回合并后图 
            results = transform(results) 
 
            if 'mix_results' in results: 
                results.pop('mix_results') 
 
        return results 

例如 Mosaic 数据增强需要一次性接收 4 张图，并输出 1 张图，则 Mosaic 类只需要实现 get_indexes 返回 4 个数据的索引，然后再调用 Mosaic 的增强函数输出 1 张大图。如果用户有其他类似需求，也只需要实现 get_indexes 和 __call__ 方法即可。

注意：get_indexes 方法能被调用的前提是你使用了 MultiImageMixDataset，经常有 issue 反应配置文件中加了 Mosaic 后没有生效，原因就是你需要同时使用 MultiImageMixDataset。

train_dataset = dict( 
    type='MultiImageMixDataset', 
    dataset=dict( 
        type=dataset_type, 
        ann_file=data_root + 'annotations/instances_train2017.json', 
        img_prefix=data_root + 'train2017/', 
        pipeline=[ 
            dict(type='LoadImageFromFile'), 
            dict(type='LoadAnnotations', with_bbox=True) 
        ], 
        filter_empty_gt=False, 
    ), 
    pipeline=train_pipeline) 
 
 
train_pipeline = [ 
    dict(type='Mosaic', img_scale=img_scale, pad_val=114.0), 
    dict( 
        type='RandomAffine', 
        scaling_ratio_range=(0.1, 2), 
        border=(-img_scale[0] // 2, -img_scale[1] // 2)), 
    dict( 
        type='MixUp', 
        img_scale=img_scale, 
        ratio_range=(0.8, 1.6), 
        pad_val=114.0), 
    dict(type='YOLOXHSVRandomAug'), 
    dict(type='RandomFlip', flip_ratio=0.5), 
    # According to the official implementation, multi-scale 
    # training is not considered here but in the 
    # 'mmdet/models/detectors/yolox.py'. 
    dict(type='Resize', img_scale=img_scale, keep_ratio=True), 
    dict( 
        type='Pad', 
        pad_to_square=True, 
        # If the image is three-channel, the pad value needs 
        # to be set separately for each channel. 
        pad_val=dict(img=(114.0, 114.0, 114.0))), 
    dict(type='FilterAnnotations', min_gt_bbox_wh=(1, 1), keep_empty=False), 
    dict(type='DefaultFormatBundle'), 
    dict(type='Collect', keys=['img', 'gt_bboxes', 'gt_labels']) 
]     

3 如何在训练中实时调整数据预处理流程以及切换 loss

这个主要是属于复现 YOLOX 算法中的需求，但是我估计有些深度用户也会有这个需求，故在本文中重点说明下当前做法。

在 YOLOX 算法中，作者采用了包括 Mosaic 、MixUp、ColorJit 等等数据增强，在 285 epoch 后要关闭 Mosaic 、MixUp 这两个数据增强，并且新增一个 L1Loss。

针对这种需求，最合理的做法是自定义相应 hook，hook 设计的初衷就是为了优雅地解决这种扩展需求。

为此我们新写了 YOLOXModeSwitchHook 类来实现上述功能。

@HOOKS.register_module() 
class YOLOXModeSwitchHook(Hook): 
 def __init__(self, 
                 num_last_epochs=15, 
                 skip_type_keys=('Mosaic', 'RandomAffine', 'MixUp')): 
        self.num_last_epochs = num_last_epochs 
        self.skip_type_keys = skip_type_keys 
        self._restart_dataloader = False 
 
    def before_train_epoch(self, runner): 
        if (epoch + 1) == runner.max_epochs - self.num_last_epochs: 
            runner.logger.info('No mosaic and mixup aug now!') 
            # The dataset pipeline cannot be updated when persistent_workers 
            # is True, so we need to force the dataloader's multi-process 
            # restart. This is a very hacky approach. 
            # 切换 pipeline 
 train_loader.dataset.update_skip_type_keys(self.skip_type_keys) 
            if hasattr(train_loader, 'persistent_workers' 
                       ) and train_loader.persistent_workers is True: 
                train_loader._DataLoader__initialized = False 
                train_loader._iterator = None 
                self._restart_dataloader = True 
            runner.logger.info('Add additional L1 loss now!') 
            # 新增 loss 
 model.bbox_head.use_l1 = True 
        else: 
            # Once the restart is complete, we need to restore 
            # the initialization flag. 
            if self._restart_dataloader: 
                train_loader._DataLoader__initialized = True 
 

上述代码还涉及到一个 DataLoader 的多进程无法修改主进行属性问题。由于这个问题相对较复杂，本文详细说明。

3.1 Dataloader 在开启多进程下无法实时修改内部属性解决办法

为了能够描述清楚这个问题，需要先说明 Dataloader 的两个重要参数：

• num_worker 开启的多进程数，如果设置为0，则只有一个主进程；如果大于1，则会开启多个子进程来加快 dataset 迭代，可以显著加快训练速度。
• persistent_workers 上一次开启的多进程是否持久化，即当前 Dataloader 迭代完后是否要回收资源，然后在下一次迭代时候重新开启。如果设置为 True，则不会回收一直复用，可以显著减少 Dataloader 切换中的耗时。

Pytorch 推荐的最佳实践是 num_worker 设置为 CPU 核心数或者 1/2，而 persistent_workers 设置为 True。

那么在上述最佳实践下，训练过程中修改 pipeline 会存在啥问题？如果你对这个问题没有啥概念，那么先看如下例子：

from torch.utils.data import Dataset, DataLoader 
import numpy as np 
 
class SimpleDataset(Dataset): 
    def __init__(self): 
        self.img_shape = (10, 10) 
    def __getitem__(self, index): 
        return np.ones(self.img_shape) 
    def __len__(self): 
        return 10 
 
def main(num_worker, persistent_workers): 
    dataset = SimpleDataset() 
    dataloader = DataLoader(dataset, num_workers=num_worker, batch_size=2, persistent_workers=persistent_workers) 
    for _ in range(2): 
        print('start epoch') 
        for i, data_batch in enumerate(dataloader): 
            print(data_batch.shape) 
            if i == 1: 
                # 在 i=1 也就是第二次迭代时候改变 shape 
                dataloader.dataset.img_shape = (20, 20) 
        print('end epoch') 
        # 在第二个 epoch 时候改变 shape 
        dataloader.dataset.img_shape = (25, 25) 
 
if __name__ == '__main__': 
    main(num_worker=2, persistent_workers=True) 
 

上述代码我们希望完成如下功能：

• 在每个 dataloader 迭代中，第 2 次迭代时候改变图片 shape 为 (20, 20)
• 在第二个 epoch 开始时候，将图片 shape 改变为 (25, 25)

在 num_worker=2, persistent_workers=True 情况下，程序运行输出为：

# num_worker=2, persistent_workers=True 
start epoch 
torch.Size([2, 10, 10]) 
torch.Size([2, 10, 10]) 
torch.Size([2, 10, 10]) 
torch.Size([2, 10, 10]) 
torch.Size([2, 10, 10]) 
end epoch 
start epoch 
torch.Size([2, 10, 10]) 
torch.Size([2, 10, 10]) 
torch.Size([2, 10, 10]) 
torch.Size([2, 10, 10]) 
torch.Size([2, 10, 10]) 
end epoch

可以发现上述两个预期一个都没有实现，这就是本文说的 Dataloader 在开启多进程下无法实时修改内部属性。

如果我们设置 num_worker=0, persistent_workers=True，即不开多进程，效果为 :：

# num_worker=0, persistent_workers=True 
ValueError: persistent_workers option needs num_workers > 0 
 

因为 persistent_workers 必须要和多进程配合使用，所以只能设置 num_worker=0, persistent_workers=False

# num_worker=0, persistent_workers=False 
start epoch 
torch.Size([2, 10, 10]) 
torch.Size([2, 10, 10]) 
torch.Size([2, 20, 20]) # 符合预期 
torch.Size([2, 20, 20]) 
torch.Size([2, 20, 20]) 
end epoch 
start epoch 
torch.Size([2, 25, 25]) # 符合预期 
torch.Size([2, 25, 25]) 
torch.Size([2, 20, 20]) # 符合预期 
torch.Size([2, 20, 20]) 
torch.Size([2, 20, 20]) 
end epoch

在 num_worker=0, persistent_workers=False 情况下发现满足了全部需求，这说明一切问题都在多进程。

那如果在 num_worker=2, persistent_workers=False 情况下会如何：

torch.Size([2, 10, 10]) 
torch.Size([2, 10, 10]) 
torch.Size([2, 10, 10]) 
torch.Size([2, 10, 10]) 
torch.Size([2, 10, 10]) 
end epoch 
start epoch 
torch.Size([2, 25, 25]) # 符合预期 
torch.Size([2, 25, 25]) 
torch.Size([2, 25, 25]) 
torch.Size([2, 25, 25]) 
torch.Size([2, 25, 25]) 
end epoch

可见 num_worker=2, persistent_workers=False 只满足了一个需求而已。总结来说是：：

• 不允许 num_worker=0, persistent_workers=True，因为 persistent_workers 要和多进程配合使用
• num_worker=0, persistent_workers=False 可以满足全部需求
• num_worker=2, persistent_workers=False 可以满足需求 2
• num_worker=2, persistent_workers=True 无法满足任何需求

从这里也可以看出，persistent_workers 只是对多进程的开启有左右，一旦多进程启动了就没啥用了，而 num_worker 直接控制了多进程的数目。

在 python 中，一旦开启多进程，那么主进程和子进程就是完全隔离的，用户无法修改任何一个进程的数据而影响其他进程的数据，除非这个数据是全局共享的。

那么在 num_worker=2, persistent_workers=True 这种情况下如何才能满足需求呢？其实需求 1 无法直接通过修改 Dataloader 参数来实现，但是需求 2 是有办法满足的。

解决办法需要从 persistent_workers 参数作用入手，说到底当其设置为 True 时候无法满足需求的原因是多进程没有重新创建，如果我们强制让他重建就可以了。

def main(num_worker, persistent_workers): 
    dataset = SimpleDataset() 
    dataloader = DataLoader(dataset, num_workers=num_worker, batch_size=2, persistent_workers=persistent_workers) 
    for _ in range(2): 
        print('start epoch') 
        for i, data_batch in enumerate(dataloader): 
            print(data_batch.shape) 
            if i == 1: 
                # 在 i=1 也就是第二次迭代时候改变 shape 
                dataloader.dataset.img_shape = (20, 20) 
        print('end epoch') 
        # 在第二个 epoch 时候改变 shape 
        dataloader.dataset.img_shape = (25, 25) 
 
 # 新增如下代码 
        if hasattr(dataloader, 'persistent_workers' 
                   ) and dataloader.persistent_workers is True: 
            dataloader._DataLoader__initialized = False 
            dataloader._iterator = None 
 

再次运行 num_worker=2, persistent_workers=True 可以得到：

start epoch 
torch.Size([2, 10, 10]) 
torch.Size([2, 10, 10]) 
torch.Size([2, 10, 10]) 
torch.Size([2, 10, 10]) 
torch.Size([2, 10, 10]) 
end epoch 
start epoch 
torch.Size([2, 25, 25]) # 符合预期 
torch.Size([2, 25, 25]) 
torch.Size([2, 25, 25]) 
torch.Size([2, 25, 25]) 
torch.Size([2, 25, 25]) 
end epoch 
 

核心就是让迭代器重建即可。在 YOLOX 中切换 pipeline 的需求就是通过上述代码实现的。

如果大家有任何疑问，欢迎留言。例如我迫切想实现需求 1，那么该如何做？ 这个问题也好解决！大家有兴趣的话，后续给安排上~

4 总结

本文重点分析了 MMDetection 中涉及到的 3 个非典型技能，主要包括：

• 如何给不同 layer 设置不同的学习率以及冻结特定层
• 如何在训练中优雅的使用多图数据增强
• 如何在训练中实时调整数据预处理流程以及切换 loss

这三个问题，我想很多人在使用 MMDetection 中碰到过的，为此本文进行了详细解答。如果你还有疑惑，可以在文章下留言，我们会积极回复补充。

以上只是我个人觉得应该重点说明的非典型操作必备技能，如果您有其他意见或者想补充的条目，欢迎留言！

在后续的文章中，我们还将带来以下技能的解读，敬请期待哦！

• 如何优雅地通过配置开启混合精度训练
• 如何在 MMDetection 中使用 timm 的骨干网络
• 为何 val workflow 的 pipeline 来自 train dataset
• DataContainer 存在的意义和作用
• 如何优雅地通过配置进行参数初始化
• 如何快速定位分布式训练中常出现的模型参数没有包括在 loss 中的错误
• EMA Hook 的正确使用方式
• 如何给 ResNet 优雅地新增插件来提升性能