训练引擎

• 训练引擎

  • 钩子(Hook)

    • 介绍

    • 默认钩子

    • MMEngine中实现的常用钩子

    • MMSelfsup 中实现的钩子

  • 优化器

    • 优化器

      • 定制 PyTorch 支持的优化器

      • 参数配置

      • MMSelfsup 中实现的优化器

    • 优化器封装

      • 梯度裁剪

      • 梯度累加

      • 自动混合精度(AMP) 训练

    • 构造器

      • MMSelfsup 中实现的构造器

钩子(Hook)

介绍

钩子机制在 OpenMMLab 开源算法库中被广泛使用,嵌入在Runner中,可以轻松管理训练过程的整个生命周期.您可以通过相关文章了解有关钩子的更多信息.

钩子只有在Runner注册后才能工作.目前钩子主要分为两类：

• 默认钩子 (default hooks)

这些钩子由Runner默认注册, 用以实现一些基本功能,并且具有默认的优先级,无需用户自行修改.

• 自定义钩子 (custom hooks)

自定义钩子通过custom_hooks注册.通常来讲,这些钩子主要用于功能增强,并且需要在配置文档中指定优先级.若没有指定钩子的优先级,则默认情况下会设置为NORMAL.

优先级列表:

Level	Value
HIGHEST	0
VERY_HIGH	10
HIGH	30
ABOVE_NORMAL	40
NORMAL(default)	50
BELOW_NORMAL	60
LOW	70
VERY_LOW	90
LOWEST	100

优先级决定了钩子的执行顺序.在训练之前，日志会打印出每个阶段的钩子的执行顺序，方便调试.

默认钩子

以下常见的钩子由 MMEngine 中的register_default_hooks 实现并在 default 中进行了注册:

Hooks	Usage	Priority
RuntimeInfoHook	将运行时间更新到消息中心.	VERY_HIGH (10)
IterTimerHook	记录迭代过程所花费的时间.	NORMAL (50)
DistSamplerSeedHook	确保分布式采样器 shuffle 是有效的.	NORMAL (50)
LoggerHook	从 Runner 的不同组件中收集日志记录, 并将其输出到终端, JSON 文件, tensorboard, wandb 等下游.	BELOW_NORMAL (60)
ParamSchedulerHook	更新优化器里面的一些超参数, 例如学习率和动量.	LOW (70)
CheckpointHook	定期保存 checkpoint.	VERY_LOW (90)

MMEngine中实现的常用钩子

在 MMEngine 中已经实现了一些钩子，它们是:

Hooks	Usage	Priority
EMAHook	在训练期间应用指数滑动平均 (Exponential Moving Average, EMA).	NORMAL (50)
EmptyCacheHook	在训练过程中释放所有未占用的缓存GPU内存.	NORMAL (50)
SyncBuffersHook	在每个Epoch结束时,在BN中同步模型缓冲区中的参数, 例如 running_mean 和 running_var.	NORMAL (50)
NaiveVisualizationHook	在测试过程中显示或写出预测结果.	LOWEST (100)

MMSelfsup 中实现的钩子

在 MMSelfsup 中已经实现了一些钩子，它们是:

• DeepClusterHook

• DenseCLHook

• ODCHook

• SimSiamHook

• SwAVHook

• ……

例如:

以 DenseCLHook 为例, 这个钩子包含了 DenseCL 中的 loss_lambda 预热.

loss_lambda 是单一和稠密对比损失的损失权重. 默认值为 0.5.

losses = dict()
losses['loss_single'] = loss_single * (1 - self.loss_lambda)
losses['loss_dense'] = loss_dense * self.loss_lambda

DenseCLHook 实现如下:

...
@HOOKS.register_module()
class DenseCLHook(Hook):
...
    def before_train_iter(self,
                          runner,
                          batch_idx: int,
                          data_batch: Optional[Sequence[dict]] = None) -> None:
...
        cur_iter = runner.iter
        if cur_iter >= self.start_iters:
            get_model(runner.model).loss_lambda = self.loss_lambda
        else:
            get_model(runner.model).loss_lambda = 0.

若钩子已在 MMEngine 或 MMSelfsup 中实现，则可以直接修改配置来使用这个钩子，如下所示:

custom_hooks = [
    dict(type='MMEngineHook', a=a_value, b=b_value, priority='NORMAL')
]

例如使用 DenseCLHook, start_iters 是 500:

custom_hooks = [
    dict(type='DenseCLHook', start_iters=500)
]

优化器

下面将通过 3 个不同的部分来介绍优化器章节: 优化器、优化器封装和构造器

优化器

定制 PyTorch 支持的优化器

我们已经支持了 PyTorch 实现的所有优化器，可参阅 mmengine/optim/optimizer/builder.py. 若要使用或修改,请更改配置文件的 optimizer 字段.

例如, 若要使用 SGD,则可进行如下修改.

optimizer = dict(type='SGD', lr=0.0003, weight_decay=0.0001)

要修改模型的学习率，只需优化器配置中修改 lr . 同时也可以根据 PyTorch 的 API doc 直接设置其他参数 .

例如,如果您期望使用如 PyTorch 中 torch.optim.Adam(params, lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0, amsgrad=False) 的 Adam 设置, 配置应该看起来像:

optimizer = dict(type='Adam', lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0, amsgrad=False)

参数配置

有些模型在优化时可能有一些特定的参数设置, 例如不需要将权重衰减应用到 BatchNorm 层和每一层的bias中. 为了精确地配置它们, 我们可以使用优化器中的 paramwise_cfg.

例如, 在 MAE 中, 我们不想对 ln, bias, pos_embed, mask_token 和 cls_token 等参数应用权重衰减, 我们可以使用以下配置文件:

optimizer = dict(
    type='AdamW', lr=1.5e-4 * 4096 / 256, betas=(0.9, 0.95), weight_decay=0.05)
optim_wrapper = dict(
    type='OptimWrapper',
    optimizer=optimizer,
    paramwise_cfg=dict(
        custom_keys={
            'ln': dict(decay_mult=0.0),
            'bias': dict(decay_mult=0.0),
            'pos_embed': dict(decay_mult=0.),
            'mask_token': dict(decay_mult=0.),
            'cls_token': dict(decay_mult=0.)
        }))

MMSelfsup 中实现的优化器

• LARS

除了PyTorch实现的优化器之外, 我们还在mmselfsup/engine/optimizers/lars.py 中实现了一个定制的 LARS,为SGD实现了分层自适应学习率缩放.

optimizer = dict(type='LARS', lr=4.8, momentum=0.9, weight_decay=1e-6)

优化器封装

除了PyTorch优化器的基本功能外, 我们还提供了一些增强功能, 例如 梯度裁剪, 梯度累加, 自动混合精度训练等. 更多细节请参考 MMEngine.

梯度裁剪

目前我们在optim_wrapper中支持clip_grad选项，您可以参考 OptimWrapper 和 PyTorch Documentation 文档了解更多的参数.下面是一个例子:

optimizer = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0001)
optim_wrapper = dict(
    type='OptimWrapper',
	optimizer=optimizer,
    clip_grad=dict(
        max_norm=0.2,
        norm_type=2))
# norm_type: type of the used p-norm, here norm_type is 2.

如果 clip_grad 不是 None, 它将是 torch.nn.utils.clip_grad.clip_grad_norm_() 的参数.

梯度累加

当没有足够的计算资源时,批量大小只能设置为一个小批量,这可能会降低模型的性能. 梯度累加可以用来解决这个问题.

下面是一个例子:

train_dataloader = dict(batch_size=64)
optim_wrapper = dict(
    type='OptimWrapper',
    optimizer=optimizer,
    accumulative_counts=4)

这个例子表示在训练期间,每4个iter进行一次反向传播. 并且上述内容相当于:

train_dataloader = dict(batch_size=256)
optim_wrapper = dict(
    type='OptimWrapper',
    optimizer=optimizer,
    accumulative_counts=1)

自动混合精度(AMP) 训练

optimizer = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0001)
optim_wrapper = dict(type='AmpOptimWrapper', optimizer=optimizer)

AmpOptimWrapper 中 loss_scale 的默认设置是 dynamic.

构造器

构造器旨在建立优化器、优化器封装以及定制不同层的超参数.配置文件中 optim_wrapper 的 paramwise_cfg 函数控制这种定制.

MMSelfsup 中实现的构造器

• LearningRateDecayOptimWrapperConstructor

LearningRateDecayOptimWrapperConstructor 为主干网络的不同层设置不同的学习率. 注意: 目前,这个优化器构造器是为 ViT, Swin, MixMIM 构建的.

一个例子:

optim_wrapper = dict(
    type='AmpOptimWrapper',
    optimizer=dict(
        type='AdamW', lr=5e-3, model_type='swin', layer_decay_rate=0.9),
    clip_grad=dict(max_norm=5.0),
    paramwise_cfg=dict(
        norm_decay_mult=0.0,
        bias_decay_mult=0.0,
        custom_keys={
            '.absolute_pos_embed': dict(decay_mult=0.0),
            '.relative_position_bias_table': dict(decay_mult=0.0)
        }),
    constructor='mmselfsup.LearningRateDecayOptimWrapperConstructor')

注意: 在 LearningRateDecayOptimWrapperConstructor 中，paramwise_cfg 只支持 weight_decay 的自定义.