Table of Contents
Shortcuts

训练引擎

MMEngine 定义了一些基础循环控制器 例如基于轮次的训练循环 (EpochBasedTrainLoop), 基于迭代次数的训练循环 (IterBasedTrainLoop), 标准的验证循环 (ValLoop) 和标准的测试循环 (TestLoop). OpenMMLab 的算法库如 MMSegmentation 将模型训练, 测试和推理抽象为执行器(Runner) 来处理. 用户可以直接使用 MMEngine 中的默认执行器, 也可以对执行器进行修改以满足定制化需求. 这个文档主要介绍用户如何配置已有的运行设定, 钩子和优化器的基本概念与使用方法.

配置运行设定

配置训练长度

循环控制器指的是训练, 验证和测试时的执行流程, 在配置文件里面使用 train_cfg, val_cfg 和 test_cfg 来构建这些流程. MMSegmentation 在 configs/_base_/schedules 文件夹里面的 train_cfg 设置常用的训练长度. 例如, 使用基于迭代次数的训练循环 (IterBasedTrainLoop) 去训练 80,000 个迭代次数, 并且每 8,000 iteration 做一次验证, 可以如下设置:

train_cfg = dict(type='IterBasedTrainLoop', max_iters=80000, val_interval=8000)

配置训练优化器

这里是一个 SGD 优化器的例子:

optim_wrapper = dict(
    type='OptimWrapper',
    optimizer=dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0005),
    clip_grad=None)

OpenMMLab 支持 PyTorch 里面所有的优化器, 更多细节可以参考 MMEngine 优化器文档.

需要强调的是, optim_wrapper 是 runner 的变量, 所以需要配置优化器时配置的字段是 optim_wrapper 字段. 更多关于优化器的使用方法, 可以看下面优化器的章节.

配置训练参数调度器

在配置训练参数调度器前, 推荐先了解 MMEngine 文档 里面关于参数调度器的基本概念.

以下是一个参数调度器的例子, 训练时前 1,000 个 iteration 时采用线性变化的学习率策略作为训练预热, 从 1,000 iteration 之后直到最后 16,000 个 iteration 时则采用默认的多项式学习率衰减:

param_scheduler = [
    dict(type='LinearLR', by_epoch=False, start_factor=0.1, begin=0, end=1000),
    dict(
        type='PolyLR',
        eta_min=1e-4,
        power=0.9,
        begin=1000,
        end=160000,
        by_epoch=False,
    )
]

注意: 当修改 train_cfg 里面 max_iters 的时候, 请确保参数调度器 param_scheduler 里面的参数也被同时修改.

钩子 (Hook)

介绍

OpenMMLab 将模型训练和测试过程抽象为 Runner, 插入钩子可以实现在 Runner 中不同的训练和测试节点 (例如 “每个训练 iter 前后”, “每个验证 iter 前后” 等不同阶段) 所需要的相应功能. 更多钩子机制的介绍可以参考这里.

Runner 中所使用的钩子分为两类:

  • 默认钩子 (default hooks)

它们实现了训练时所必需的功能, 在配置文件中用 default_hooks 定义传给 Runner, Runner 通过 register_default_hooks 方法注册. 钩子有对应的优先级, 优先级越高, 越早被执行器调用. 如果优先级一样, 被调用的顺序和钩子注册的顺序一致. 不建议用户修改默认钩子的优先级, 可以参考 mmengine hooks 文档 了解钩子优先级的定义. 下面是 MMSegmentation 中所用到的默认钩子：

钩子 功能 优先级
IterTimerHook 记录 iteration 花费的时间. NORMAL (50)
LoggerHook 从 Runner 里不同的组件中收集日志记录, 并将其输出到终端, JSON 文件, tensorboard, wandb 等下游. BELOW_NORMAL (60)
ParamSchedulerHook 更新优化器里面的一些超参数, 例如学习率的动量. LOW (70)
CheckpointHook 规律性地保存 checkpoint 文件. VERY_LOW (90)
DistSamplerSeedHook 确保分布式采样器 shuffle 是打开的. NORMAL (50)
SegVisualizationHook 可视化验证和测试过程里的预测结果. NORMAL (50)

MMSegmentation 会在 defualt_hooks 里面注册一些训练所必需功能的钩子::

default_hooks = dict(
    timer=dict(type='IterTimerHook'),
    logger=dict(type='LoggerHook', interval=50, log_metric_by_epoch=False),
    param_scheduler=dict(type='ParamSchedulerHook'),
    checkpoint=dict(type='CheckpointHook', by_epoch=False, interval=32000),
    sampler_seed=dict(type='DistSamplerSeedHook'),
    visualization=dict(type='SegVisualizationHook'))

以上默认钩子除 SegVisualizationHook 外都是在 MMEngine 中所实现, SegVisualizationHook 是在 MMSegmentation 里被实现的钩子, 之后会专门介绍.

  • 修改默认的钩子

以 default_hooks 里面的 logger 和 checkpoint 为例, 我们来介绍如何修改 default_hooks 中默认的钩子.

(1) 模型保存配置

default_hooks 使用 checkpoint 字段来初始化模型保存钩子 (CheckpointHook).

checkpoint = dict(type='CheckpointHook', interval=1)

用户可以设置 max_keep_ckpts 来只保存少量的检查点或者用 save_optimizer 来决定是否保存 optimizer 的信息. 更多相关参数的细节可以参考这里.

(2) 日志配置

日志钩子 (LoggerHook) 被用来收集 执行器 (Runner) 里面不同组件的日志信息然后写入终端, JSON 文件, tensorboard 和 wandb 等地方.

logger=dict(type='LoggerHook', interval=10)

在最新的 1.x 版本的 MMSegmentation 里面, 一些日志钩子 (LoggerHook) 例如 TextLoggerHook, WandbLoggerHook 和 TensorboardLoggerHook 将不再被使用. 作为替代, MMEngine 使用 LogProcessor 来处理上述钩子处理的信息, 它们现在在 MessageHub, WandbVisBackend 和 TensorboardVisBackend 里面.

具体使用方法如下, 配置可视化器和同时指定可视化后端, 这里使用 Tensorboard 作为可视化器的后端:

# TensorboardVisBackend
visualizer = dict(
    type='SegLocalVisualizer', vis_backends=[dict(type='TensorboardVisBackend')], name='visualizer')

关于更多相关用法, 可以参考 MMEngine 可视化后端用户教程.

  • 自定义钩子 (custom hooks)

自定义钩子在配置通过 custom_hooks 定义, Runner 通过 register_custom_hooks 方法注册. 自定义钩子优先级需要在配置文件里设置, 如果没有设置, 则会被默认设置为 NORMAL. 下面是部分 MMEngine 中实现的自定义钩子:

钩子 用法
EMAHook 在模型训练时使用指数滑动平均 (Exponential Moving Average, EMA).
EmptyCacheHook 在训练时释放所有没有被缓存占用的 GPU 显存.
SyncBuffersHook 在每个训练 Epoch 结束时同步模型 buffer 里的参数例如 BN 里的 running_mean 和 running_var.

以下是 EMAHook 的用例, 配置文件中, 将已经实现的自定义钩子的配置作为 custom_hooks 列表中的成员.

custom_hooks = [
    dict(type='EMAHook', start_iters=500, priority='NORMAL')
]

SegVisualizationHook

MMSegmentation 实现了 SegVisualizationHook, 用来在验证和测试时可视化预测结果. SegVisualizationHook 重写了基类 Hook 中的 _after_iter 方法, 在验证或测试时, 根据指定的迭代次数间隔调用 visualizer 的 add_datasample 方法绘制语义分割结果, 具体实现如下:

...
@HOOKS.register_module()
class SegVisualizationHook(Hook):
...
    def _after_iter(self,
                    runner: Runner,
                    batch_idx: int,
                    data_batch: dict,
                    outputs: Sequence[SegDataSample],
                    mode: str = 'val') -> None:
...
        # 如果是训练阶段或者 self.draw 为 False 则直接跳出
        if self.draw is False or mode == 'train':
            return
...
        if self.every_n_inner_iters(batch_idx, self.interval):
            for output in outputs:
                img_path = output.img_path
                img_bytes = self.file_client.get(img_path)
                img = mmcv.imfrombytes(img_bytes, channel_order='rgb')
                window_name = f'{mode}_{osp.basename(img_path)}'

                self._visualizer.add_datasample(
                    window_name,
                    img,
                    data_sample=output,
                    show=self.show,
                    wait_time=self.wait_time,
                    step=runner.iter)

关于可视化更多的细节可以查看这里.

优化器

在上面配置运行设定里, 我们给出了配置训练优化器的简单示例. 本章节将进一步详细介绍在 MMSegmentation 里如何配置优化器.

优化器封装

OpenMMLab 2.0 设计了优化器封装, 它支持不同的训练策略, 包括混合精度训练、梯度累加和梯度截断等, 用户可以根据需求选择合适的训练策略. 优化器封装还定义了一套标准的参数更新流程, 用户可以基于这一套流程, 在同一套代码里, 实现不同训练策略的切换. 如果想了解更多, 可以参考 MMEngine 优化器封装文档.

以下是 MMSegmentation 中常用的使用方法:

配置 PyTorch 支持的优化器

OpenMMLab 2.0 支持 PyTorch 原生所有优化器, 参考这里.

在配置文件中设置训练时 Runner 所使用的优化器, 需要定义 optim_wrapper, 而不是 optimizer, 下面是一个配置训练中优化器的例子:

optim_wrapper = dict(
    type='OptimWrapper',
    optimizer=dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0005),
    clip_grad=None)

配置梯度裁剪

当模型训练需要使用梯度裁剪的训练技巧式, 可以按照如下示例进行配置:

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.01, norm_type=2))

这里 max_norm 指的是裁剪后梯度的最大值, norm_type 指的是裁剪梯度时使用的范数. 相关方法可参考 torch.nn.utils.clip_grad_norm_.

配置混合精度训练

当需要使用混合精度训练降低内存时, 可以使用 AmpOptimWrapper, 具体配置如下:

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 来实现.

下面的配置文件以 ViT optim_wrapper 为例介绍 paramwise_cfg 参数使用. 训练时将 pos_embed, mask_token, norm 模块的 weight decay 参数的系数设置成 0. 即: 在训练时, 这些模块的 weight decay 将被变为 weight_decay * decay_mult=0.

optimizer = dict(
        type='AdamW', lr=0.00006, betas=(0.9, 0.999), weight_decay=0.01)
optim_wrapper = dict(
    type='OptimWrapper',
    optimizer=optimizer,
    paramwise_cfg=dict(
        custom_keys={
            'pos_embed': dict(decay_mult=0.),
            'cls_token': dict(decay_mult=0.),
            'norm': dict(decay_mult=0.)
        }))

其中 decay_mult 指的是对应参数的权重衰减的系数. 关于更多 paramwise_cfg 的使用可以在 MMEngine 优化器封装文档 里面查到.

优化器封装构造器

默认的优化器封装构造器 DefaultOptimWrapperConstructor 根据输入的 optim_wrapper 和 optim_wrapper 中定义的 paramwise_cfg 来构建训练中使用的优化器. 当 DefaultOptimWrapperConstructor 功能不能满足需求时, 可以自定义优化器封装构造器来实现超参数的配置.

MMSegmentation 中的实现了 LearningRateDecayOptimizerConstructor, 可以对以 ConvNeXt, BEiT 和 MAE 为骨干网络的模型训练时, 骨干网络的模型参数的学习率按照定义的衰减比例（decay_rate）逐层递减, 在配置文件中的配置如下:

optim_wrapper = dict(
    _delete_=True,
    type='AmpOptimWrapper',
    optimizer=dict(
        type='AdamW', lr=0.0001, betas=(0.9, 0.999), weight_decay=0.05),
    paramwise_cfg={
        'decay_rate': 0.9,
        'decay_type': 'stage_wise',
        'num_layers': 12
    },
    constructor='LearningRateDecayOptimizerConstructor',
    loss_scale='dynamic')

_delete_=True 的作用是 OpenMMLab Config 中的忽略继承的配置, 在该代码片段中忽略继承的 optim_wrapper 配置, 更多 _delete_ 字段的内容可以参考 MMEngine 文档.

Read the Docs v: main
Versions
latest
0.x
main
dev-1.x
Downloads
epub
On Read the Docs
Project Home
Builds
Free document hosting provided by Read the Docs.