过节福利 | MMCV Hook 超全使用方法（下）-阿里云开发者社区

OptimizerHook

介绍

OptimizerHook 包含一些 optimizer 相关的操作：

- 梯度清零 runner.optimizer.zero_grad()

- 反向传播 runner.output['loss'].backward()

- 梯度阶段 clip_grads（可选）

- 参数更新 runner.optimizer.step()

MMCV 还提供了 Fp16OptimizerHook 和 GradientCumulativeOptimizerHook，前者用于混合精度训练，后者用于梯度累计。

Fp16OptimizerHook 是混合精度训练在 MMCV 中的实现，主要逻辑如下：

- 维护一个 FP32 数值精度模型的副本

- 在每个 iteration

拷贝并且转换成 FP16 模型

前向传播（FP16 的模型参数)，此时 weights, activations 都是 FP16

loss 乘 scale factor s

反向传播（FP16 的模型参数和参数梯度)，此时 gradients 也是 FP16

参数梯度乘 1/s

利用 FP16 的梯度更新 FP32 的模型参数

GradientCumulativeOptimizerHook 用于节省显存，即通过指定梯度累积的次数，实现反向传播多次才更新参数，常常用于显存不足但想用比较大的 batch size 训练模型。

用法

- 最简用法

optimizer_config = dict(grad_clip=None)
runner = EpochBasedRunner(...)
runner.register_optimizer_hook(optimizer_config)

- 梯度截断

用于避免梯度爆炸，grad_clip 的设置可参考：

https://pytorch.org/docs/stable/generated/torch.nn.utils.clip_grad_norm_.html

optimizer_config=dict(grad_clip=dict(max_norm=35, norm_type=2))
runner = EpochBasedRunner(...)
runner.register_optimizer_hook(optimizer_config)

- 检测异常的模型参数（mmcv>=1.4.1）

该设置会降低训练速度，故只应用于调试。该参数用于寻找不参与计算图的模型参数，不参与计算图的模型参数包括两种情况，一种是该模型参数没有参与前向计算，另一种参与了前向计算但没有参与 loss 的计算。

optimizer_config = dict(grad_clip=None, detect_anomalous_params=True)
runner = EpochBasedRunner(...)
runner.register_optimizer_hook(optimizer_config)

- 梯度累积

可用于显存不足的情况下训练训练更大的 batch size，更多细节可参考：

https://github.com/open-mmlab/mmcv/pull/1221

optimizer_config = dict(type="GradientCumulativeOptimizerHook", cumulative_iters=4)
runner = EpochBasedRunner(...)
runner.register_optimizer_hook(optimizer_config)

！! 切记我们不能同时使用 OptimizerHook 和 GradientCumulativeOptimizerHook，否则会遇到运行时错误: Trying to backward through the graph a second time，更多细节见：

https://github.com/open-mmlab/mmcv/issues/1379

- 混合精度训练

使用 MMCV 的 AMP 功能，只需遵循以下几个步骤：

1. 将 auto_fp16 装饰器应用到 model 的 forward 函数上

2. 设置模型的 fp16_enabled 为 True 表示开启 AMP 训练，否则不生效

3. 如果开启了 AMP，需要同时配置对应的 FP16 优化器配置 Fp16OptimizerHook

4. 在训练的不同时刻，调用 Fp16OptimizerHook，如果你同时使用了 MMCV 中的 Runner 模块，那么直接将第 3 步的参数输入到 Runner 中即可

5. (可选) 如果对应某些 OP 希望强制运行在 FP32 上，则可以在对应位置引入 force_fp32 装饰器

# 1 作用到 forward 函数中
class ExampleModule(nn.Module):
    @auto_fp16()
    def forward(self, x, y):
        return x, y
# 2 如果开启 AMP，则需要加入开启标志
model.fp16_enabled = True     
# 3 配置 Fp16OptimizerHook
optimizer_config = Fp16OptimizerHook(
    **cfg.optimizer_config, **fp16_cfg, distributed=distributed)
# 4 传递给 runner
runner.register_training_hooks(cfg.lr_config, optimizer_config,
                               cfg.checkpoint_config, cfg.log_config,
                               cfg.get('momentum_config', None))   
# 5 可选
class ExampleModule(nn.Module):
    @auto_fp16()
    def forward(self, x, y):
        features=self._forward(x, y)
        loss=self._loss(features,labels)
        return loss
    def _forward(self, x, y):
       pass
    @force_fp32(apply_to=('features',))
    def _loss(features,labels) :
        pass

！注意：force_fp32 要生效，依然需要 fp16_enabled 为 True

EMAHook

介绍

EMAHook 使用滑动平均策略对模型参数做平均以提高训练过程的稳定性。

EMAHook 链接：

https://github.com/open-mmlab/mmcv/blob/master/mmcv/runner/hooks/ema.py

！需保证 EMAHook 的优先级高于 CheckpointHook

用法

- 最简用法

ema_config = dict(type="EMAHook")
runner = EpochBasedRunner(...)
runner.register_hook(ema_config, priority='NORMAL')

- 设置更新 ema 参数的间隔数

默认是每一次迭代之后都更新 ema 参数，我们可以通过 interval 设置间隔

ema_config = dict(type="EMAHook", interval=4)
runner = EpochBasedRunner(...)
runner.register_hook(ema_config, priority='NORMAL')

- 恢复训练

当中断训练之后恢复训练，我们需要设置 resume_from 参数（注意，这个参数是一定要设置的，否则和不中断训练的结果是不一致的），这样我们就可以加载 ema 参数，以保证使用 EMAHook 训练的正确性。

ema_config = dict(type="EMAHook", resume_from="/path/of/your/checkpoint.pth")
runner = EpochBasedRunner(...)
runner.register_hook(ema_config, priority='NORMAL')

LrUpdaterHook

介绍

学习率决定每次更新的步长，合适的学习率可以使训练快速收敛。MMCV 中提供很多学习率衰减策略，其中部分学习率衰减策略也伴有动量衰减策略。

下表是 MMCV 提供的学习率更新策略和对应的动量更新策略——

学习率调整策略虽然有很多种，但用法如出一辙，下面只举两个常用的用法。

用法

- 等间隔调整学习率

lr_config = dict(
    policy='step',  # 对应 StepLrUpdaterHook
    step=[16, 19])
runner = EpochBasedRunner(...)
runner.register_lr_hook(lr_config)

- 余弦退火调整学习率

lr_config = dict(
    policy='CosineAnnealing',  # 对应 CosineAnnealingLrUpdaterHook
    min_lr=0.01,
    )
runner = EpochBasedRunner(...)
runner.register_lr_hook(lr_config)

LoggerHook

MMCV 提供以下日志相关的 LoggerHook

面对这么多 LoggerHook，我们该如何选择合适的日志 Hook 呢？以下是我们的一些推荐：

- 如果你偏好在本地使用，想要一款能满足基本的实验记录需求、且上手容易的，推荐使用 TensorBoard；

- 如果对可视化和实验记录有较高要求，推荐使用 Neptune 和 WandB 。两者的实验记录功能都非常强大，并且都支持多人协作，适合大型项目的管理；

- 如果你是开源项目的爱好者，或者你希望记录完整的 ML 实验全过程，MLflow 是一个不错的选择；

- 如果你只需要记录基本的实验数据，但是对于大规模数据的版本管理有比较高的需求，那么推荐你使用 DVC。

事实上，只要你喜欢，你可以同时使用的所有 LoggerHook。

! 注意，MMCV 目前提供的 LoggerHook 功能有限，只用于记录标量数据，例如 loss、acc 等，暂不支持可视化图像或特征图等。欢迎提 PR 完善 LoggerHook。

TextLoggerHook

介绍

TextLoggerHook 会将日志打印到终端以及保存到 json 文件。

用法

- 最简用法

间隔 100 个 iteration 打印一次日志

log_config = {
    'interval': 100,
    'hooks': [
        {
            'type': 'TextLoggerHook',
        },
     ]
}
runner = EpochBasedRunner(...)
runner.register_logger_hooks(log_config)

- 训练完成后将日志拷贝至指定路径

log_config = {
    'hooks': [
        {
            'type': 'TextLoggerHook',
            'out_dir': '/path/of/expected_directory',
            'out_suffix': ('.log.json', '.log'),  # 可以指定拷贝的文件后缀
        },
     ]
}
runner = EpochBasedRunner(...)
runner.register_logger_hooks(log_config)

TensorBoard

介绍

TensorBoard 最初是随 TensorFlow 提出的一款可视化工具包，其便捷性和完善的记录功能使它得到了广泛应用，并扩展到 PyTorch 等多种深度学习框架。

TensorBoard 支持记录多种数据类型：

- 指标和损失

- 超参数和模型 config

- 图片数据（可视化权重、张量、多个图像）

- 模型图

- Embedding Projector（在低维空间可视化高维数据）

！注意：MMCV 中提供的 TensorboardLoggerHook 只支持记录指标和损失。

安装

pip install tensorboard

用法

间隔 100 个 iteration 往 TensorBoard 写一次日志

log_config = {
    'interval': 100,
    'hooks': [
        {
            'type': 'TensorboardLoggerHook',
        },
     ]
}

将 https://github.com/open-mmlab/mmcv/blob/master/examples/train.py

中的 log_config 替换为上述配置，然后启动训练

python examples/train.py

另起一个终端输入

tensorboard --logdir work_dirs
# 也可以将 log 上传至 TensorBoard.dev
tensorboard dev upload --logdir work_dirs

打开 chrome 浏览器，输入 http://localhost:6006/ （可点击 shared_logs 查看我共享的日志）

Neptune

介绍

Neptune 是一个集实验记录、数据存储、可视化、模型注册等多种功能于一体的机器学习实验管理工具，用户可以在网页端轻松地查看所有的记录与可视化结果。Neptune 支持记录的数据类型包括但不限于：

- 指标和损失

- 超参数和模型 config

- 模型 checkpoints

- Git 信息

- 数据版本管理

- 硬件消耗

- 文件

- 控制台日志

- 图片数据（图片文件、Matplotlib figure、PIL image、Numpy array、Tensor）

- 交互式可视化（自动将 Matplotlib figure 转为交互式，同时支持其他格式如 html 文件、Altair chart）

相较于 TensorBoard，Neptune 支持记录更多种类的数据，并且提供了用户友好的 UI，使用户可以灵活地调整可视化界面。Neptune 还提供了 TensorBoard 接口，可以很方便地把 TensorBoard logs 转换为 Neptune experiments。

！注意：MMCV 中提供的 NeptuneLoggerHook 只支持记录指标和损失

安装

- 安装 Neptune

pip install neptune-client

- 注册 neptune 账号并设置 NEPTUNE_API_TOKEN

neptune 账号链接：http://neptune.ai

NEPTUNE_API_TOKEN 链接：

https://docs.neptune.ai/getting-started/installation#authentication-neptune-api-token

用法

间隔 100 个 iteration 往 Neptune 写一次日志

log_config = {
     'hooks': [
        {
            'type': 'NeptuneLoggerHook',
            'init_kwargs': {
                # YOUR_WORKSPACE 是账号名，YOUR_PROJECT 是项目名
                'project': '<YOUR_WORKSPACE/YOUR_PROJECT>',
            },
        },
     ]
}

将 https://github.com/open-mmlab/mmcv/blob/master/examples/train.py 中的 log_config 替换为上述配置，然后启动训练

python examples/train.py

打开 https://neptune.ai/ 并登录即可查看日志（可点击 shared_logs 查看我共享的日志）

如果我们实现一个定制化的 Hook，我们要做的是考虑需要在 Hook 的哪些方法中添加逻辑。

例如，我们想在训练的过程中判断 loss 是否有效（无穷大即为无效），我们可以在每次迭代之后判断 loss 的值，即可以在 after_train_iter 中添加判断的逻辑。

！注意：如无必要，不应当在 Hook 中修改能够影响其他 Hook 的属性或者方法。而且原则上 Hook 之间最好不要有前后依赖关系。Hook 的主要目的是扩展功能，而不是修改已经实现的功能。

# https://github.com/open-mmlab/mmdetection/blob/master/mmdet/core/hook/checkloss_hook.py
import torch
from mmcv.runner.hooks import HOOKS, Hook
@HOOKS.register_module()
class CheckInvalidLossHook(Hook):
    """Check invalid loss hook.
    This hook will regularly check whether the loss is valid
    during training.
    Args:
        interval (int): Checking interval (every k iterations).
            Default: 50.
    """
    def __init__(self, interval=50):
        self.interval = interval
    def after_train_iter(self, runner):
        if self.every_n_iters(runner, self.interval):
            assert torch.isfinite(runner.outputs['loss']), \
                runner.logger.info('loss become infinite or NaN!')

文章来源：公众号【OpenMMLab】

2021-12-24 13:59

过节福利 | MMCV Hook 超全使用方法（下）

OptimizerHook

- 梯度截断

- 梯度累积

EMAHook

LrUpdaterHook

LoggerHook

TextLoggerHook

TensorBoard

Neptune

热门文章

最新文章

相关电子书