前言✦
大家好,今天我们将开启全新的 MMDetection 系列文章,是时候带大家学习一些非典型操作技能啦。
这些非典型操作出现的原因各种各样,有部分来自内部和社区用户所提需求,有部分来自复现算法本身的需求。希望大家通过学习本系列文章,在使用 MMDetection 进行扩展开发时可以更加游刃有余,轻松秀出各种骚操作。
本文是非典型操作系列文章的首篇,所涉及到的典型操作技能为:
如何给不同 layer 设置不同的学习率以及冻结特定层
如何在训练中优雅地使用多图数据增强
如何在训练中实时调整数据预处理流程以及切换 loss
注意!
1. 本文需要用户对 MMDetection 本身有一定了解,可通过官方文档了解 MMDetection。
2. 本文所述非典型操作办法可能仅适用于 MMDetection V2.21.0 及其以前版本,随着 MMDetection 持续更新,相信之后会有更优雅的解决办法。
MMDetection 官方文档:
https://github.com/open-mmlab/mmdetection/tree/master/docs/zh_cn
如何给不同 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,从而间接达到目标。
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
冻结特定层解决办法
对于没有 BN 层的模块,用户可以将想冻结的层中的 lr_mult 设置为 0。但是一旦有 BN,lr=0 只是可学习参数不再更新,但是全局均值和方差依然在改,没有实现真正的冻结。
如果面对有 BN 层的冻结需求时,暂时没有办法通过直接修改配置实现,目前来看有两种自定义办法:
- 自定义 OptimizerConstructor 或者继承 DefaultOptimizerConstructor,然后在内部自己处理逻辑
- 直接在构建模型的时候将想要冻结的层设置 requires_grad 属性为 False,并且切换为 eval 模式
对于一般水平用户推荐第二种最简单直接的做法,如果是有能力自定义 OptimizerConstructor 的用户则推荐直接自定义,这样写更加通用。
如何在训练中优雅地使用多张图数据增强
同时使用多张图数据增强的典型代表是 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']) ]
如何在训练中实时调整数据预处理流程以及切换 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 的多进程无法修改主进行属性问题。由于这个问题相对较复杂,本文详细说明。
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,那么该如何做?这个问题也好解决!大家有兴趣的话,后续给安排上~
总结✦
本文重点分析了 MMDetection 中涉及到的 3 个非典型技能,主要包括:
如何给不同 layer 设置不同的学习率以及冻结特定层
如何在训练中优雅地使用多图数据增强
如何在训练中实时调整数据预处理流程以及切换 loss
这三个问题,我想很多人在使用 MMDetection 中碰到过的,为此本文进行了详细解答。如果你还有疑惑,可以在文章下留言,我们会积极回复补充。
文章来源:【OpenMMLab】
2022-02-28 18:08