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
