自定义训练优化策略¶

在我们的算法库中，已经提供了通用数据集（如ImageNet，CIFAR）的默认训练策略配置。如果想要在这些数据集上继续提升模型性能，或者在不同数据集和方法上进行新的尝试，我们通常需要修改这些默认的策略。

在本教程中，我们将介绍如何在运行自定义训练时，通过修改配置文件进行构造优化器、参数化精细配置、梯度裁剪、梯度累计以及定制动量调整策略等。同时也会通过模板简单介绍如何自定义开发优化器和构造器。

配置训练优化策略¶

我们通过 optim_wrapper 来配置主要的优化策略，包括优化器的选择，混合精度训练的选择，参数化精细配置，梯度裁剪以及梯度累计。接下来将分别介绍这些内容。

构造 PyTorch 内置优化器¶

MMPretrain 支持 PyTorch 实现的所有优化器，仅需在配置文件中，指定优化器封装需要的 optimizer 字段。

如果要使用 SGD，则修改如下。这里要注意所有优化相关的配置都需要封装在 optim_wrapper 配置里。

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

备注

配置文件中的 ‘type’ 不是构造时的参数，而是 PyTorch 内置优化器的类名。更多优化器选择可以参考PyTorch 支持的优化器列表。

要修改模型的学习率，只需要在优化器的配置中修改 lr 即可。要配置其他参数，可直接根据 PyTorch API 文档进行。

例如，如果想使用 Adam 并设置参数为 torch.optim.Adam(params, lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0, amsgrad=False)。则需要进行如下修改：

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

备注

考虑到对于单精度训练来说，优化器封装的默认类型就是 OptimWrapper，我们在这里可以直接省略，因此配置文件可以进一步简化为：

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

混合精度训练¶

如果我们想要使用混合精度训练（Automactic Mixed Precision），我们只需简单地将 optim_wrapper 的类型改为 AmpOptimWrapper。

optim_wrapper = dict(type='AmpOptimWrapper', optimizer=...)

另外，为了方便，我们同时在启动训练脚本 tools/train.py 中提供了 --amp 参数作为开启混合精度训练的开关，更多细节可以参考训练教程。

参数化精细配置¶

在一些模型中，不同的优化策略需要适应特定的参数，例如不在 BatchNorm 层使用权重衰减，或者在不同层使用不同的学习率等等。我们需要用到 optim_wrapper 中的 paramwise_cfg 参数来进行精细化配置。

为不同类型的参数设置超参乘子

例如，我们可以在 paramwise_cfg 配置中设置 norm_decay_mult=0. 来改变归一化层权重和偏移的衰减为0。
```
optim_wrapper = dict(
    optimizer=dict(type='SGD', lr=0.8, weight_decay=1e-4),
    paramwise_cfg=dict(norm_decay_mult=0.))
```
支持更多类型的参数配置，参考以下列表：
- bias_lr_mult：偏置的学习率系数（不包括正则化层的偏置以及可变形卷积的 offset），默认值为 1
- bias_decay_mult：偏置的权值衰减系数（不包括正则化层的偏置以及可变形卷积的 offset），默认值为 1
- norm_decay_mult：正则化层权重和偏置的权值衰减系数，默认值为 1
- flat_decay_mult: 一维参数的权值衰减系数，默认值为 1
- dwconv_decay_mult：Depth-wise 卷积的权值衰减系数，默认值为 1
- bypass_duplicate：是否跳过重复的参数，默认为 False
- dcn_offset_lr_mult：可变形卷积（Deformable Convolution）的学习率系数，默认值为 1

为特定参数设置超参乘子

MMPretrain 通过 paramwise_cfg 的 custom_keys 参数来配置特定参数的超参乘子。

例如，我们可以通过以下配置来设置所有 backbone.layer0 层的学习率和权重衰减为0， backbone 的其余层和优化器保持一致，另外 head 层的学习率为0.001.

optim_wrapper = dict(
    optimizer=dict(type='SGD', lr=0.01, weight_decay=0.0001),
    paramwise_cfg=dict(
        custom_keys={
            'backbone.layer0': dict(lr_mult=0, decay_mult=0),
            'backbone': dict(lr_mult=1),
            'head': dict(lr_mult=0.1)
        }))

梯度裁剪¶

在训练过程中，损失函数可能接近于一些异常陡峭的区域，从而导致梯度爆炸。而梯度裁剪可以帮助稳定训练过程，更多介绍可以参见该页面。

目前我们支持在 optim_wrapper 字段中添加 clip_grad 参数来进行梯度裁剪，更详细的参数可参考 PyTorch 文档。

用例如下：

optim_wrapper = dict(
    optimizer=dict(type='SGD', lr=0.01, weight_decay=0.0001),
    # norm_type: 使用的范数类型，此处使用范数2。
    clip_grad=dict(max_norm=35, norm_type=2))

梯度累计¶

计算资源缺乏缺乏时，每个训练批次的大小（batch size）只能设置为较小的值，这可能会影响模型的性能。

可以使用梯度累计来规避这一问题。我们支持在 optim_wrapper 字段中添加 accumulative_counts 参数来进行梯度累计。

用例如下：

train_dataloader = dict(batch_size=64)
optim_wrapper = dict(
    optimizer=dict(type='SGD', lr=0.01, weight_decay=0.0001),
    accumulative_counts=4)

表示训练时，每 4 个 iter 执行一次反向传播。由于此时单张 GPU 上的批次大小为 64，也就等价于单张 GPU 上一次迭代的批次大小为 256，也即：

train_dataloader = dict(batch_size=256)
optim_wrapper = dict(
    optimizer=dict(type='SGD', lr=0.01, weight_decay=0.0001))

配置参数优化策略¶

在训练过程中，优化参数例如学习率、动量，通常不会是固定不变，而是随着训练进程的变化而调整。PyTorch 支持一些学习率调整的调度器，但是不足以完成复杂的策略。在 MMPretrain 中，我们提供 param_scheduler 来更好地控制不同优化参数的策略。

配置学习率调整策略¶

深度学习研究中，广泛应用学习率衰减来提高网络的性能。我们支持大多数 PyTorch 学习率调度器，其中包括 ExponentialLR, LinearLR, StepLR, MultiStepLR 等等。

单个学习率策略

多数情况下，我们使用单一学习率策略，这里 param_scheduler 会是一个字典。比如在默认的 ResNet 网络训练中，我们使用阶梯式的学习率衰减策略 MultiStepLR，配置文件为：
```
param_scheduler = dict(
    type='MultiStepLR',
    by_epoch=True,
    milestones=[100, 150],
    gamma=0.1)
```
或者我们想使用 CosineAnnealingLR 来进行学习率衰减：
```
param_scheduler = dict(
    type='CosineAnnealingLR',
    by_epoch=True,
    T_max=num_epochs)
```
多个学习率策略

然而在一些其他情况下，为了提高模型的精度，通常会使用多种学习率策略。例如，在训练的早期阶段，网络容易不稳定，而学习率的预热就是为了减少这种不稳定性。

整个学习过程中，学习率将会通过预热从一个很小的值逐步提高到预定值，再会通过其他的策略进一步调整。

在 MMPretrain 中，我们同样使用 param_scheduler ，将多种学习策略写成列表就可以完成上述预热策略的组合。

例如：
1. 在前50次迭代中逐迭代次数地线性预热
```
  param_scheduler = [
      # 逐迭代次数，线性预热
      dict(type='LinearLR',
          start_factor=0.001,
          by_epoch=False,  # 逐迭代次数
          end=50),  # 只预热50次迭代次数
      # 主要的学习率策略
      dict(type='MultiStepLR',
          by_epoch=True,
          milestones=[8, 11],
          gamma=0.1)
  ]
```
1. 在前10轮迭代中逐迭代次数地线性预热
```
  param_scheduler = [
      # 在前10轮迭代中，逐迭代次数，线性预热
      dict(type='LinearLR',
          start_factor=0.001,
          by_epoch=True,
          end=10,
          convert_to_iter_based=True,  # 逐迭代次数更新学习率.
      ),
      # 在 10 轮次后，通过余弦退火衰减
      dict(type='CosineAnnealingLR', by_epoch=True, begin=10)
  ]
```
注意这里增加了 begin 和 end 参数，这两个参数指定了调度器的生效区间。生效区间通常只在多个调度器组合时才需要去设置，使用单个调度器时可以忽略。当指定了 begin 和 end 参数时，表示该调度器只在 [begin, end) 区间内生效，其单位是由 by_epoch 参数决定。在组合不同调度器时，各调度器的 by_epoch 参数不必相同。如果没有指定的情况下，begin 为 0， end 为最大迭代轮次或者最大迭代次数。

如果相邻两个调度器的生效区间没有紧邻，而是有一段区间没有被覆盖，那么这段区间的学习率维持不变。而如果两个调度器的生效区间发生了重叠，则对多组调度器叠加使用，学习率的调整会按照调度器配置文件中的顺序触发（行为与 PyTorch 中 ChainedScheduler 一致）。

小技巧

为了避免学习率曲线与预期不符，配置完成后，可以使用 MMPretrain 提供的学习率可视化工具画出对应学习率调整曲线。

配置动量调整策略¶

MMPretrain 支持动量调度器根据学习率修改优化器的动量，从而使损失函数收敛更快。用法和学习率调度器一致。

我们支持的动量策略和详细的使用细节可以参考这里。我们只将调度器中的 LR 替换为了 Momentum，动量策略可以直接追加 param_scheduler 列表中。

这里是一个用例：

param_scheduler = [
    # 学习率策略
    dict(type='LinearLR', ...),
    # 动量策略
    dict(type='LinearMomentum',
         start_factor=0.001,
         by_epoch=False,
         begin=0,
         end=1000)
]

新增优化器或者优化器构造器¶

备注

本部分将修改 MMPretrain 源码或者向 MMPretrain 框架添加代码，初学者可跳过。

新增优化器¶

在学术研究和工业实践中，可能需要使用 MMPretrain 未实现的优化方法，可以通过以下方法添加。

定义一个新的优化器

一个自定义的优化器可根据如下规则进行定制：

假设我们想添加一个名为 MyOptimzer 的优化器，其拥有参数 a, b 和 c。可以创建一个名为 mmpretrain/engine/optimizer 的文件夹，并在目录下的一个文件，如 mmpretrain/engine/optimizer/my_optimizer.py 中实现该自定义优化器：
```
from mmpretrain.registry import OPTIMIZERS
from torch.optim import Optimizer


@OPTIMIZERS.register_module()
class MyOptimizer(Optimizer):

    def __init__(self, a, b, c):
        ...

    def step(self, closure=None):
        ...
```
注册优化器

要注册上面定义的上述模块，首先需要将此模块导入到主命名空间中。有两种方法可以实现它。

修改 mmpretrain/engine/optimizers/__init__.py，将其导入至 mmpretrain.engine 包。
```
# 在 mmpretrain/engine/optimizers/__init__.py 中
...
from .my_optimizer import MyOptimizer # MyOptimizer 是我们自定义的优化器的名字

__all__ = [..., 'MyOptimizer']
```
在运行过程中，我们会自动导入 mmpretrain.engine 包并同时注册 MyOptimizer。
在配置文件中指定优化器

之后，用户便可在配置文件的 optim_wrapper.optimizer 域中使用 MyOptimizer：
```
optim_wrapper = dict(
    optimizer=dict(type='MyOptimizer', a=a_value, b=b_value, c=c_value))
```

新增优化器构造器¶

某些模型可能具有一些特定于参数的设置以进行优化，例如为所有 BatchNorm 层设置不同的权重衰减。

尽管我们已经可以使用 optim_wrapper.paramwise_cfg 字段来配置特定参数的优化设置，但可能仍然无法覆盖你的需求。

当然你可以在此基础上进行修改。我们默认使用 DefaultOptimWrapperConstructor 来构造优化器。在构造过程中，通过 paramwise_cfg 来精细化配置不同设置。这个默认构造器可以作为新优化器构造器实现的模板。

我们可以新增一个优化器构造器来覆盖这些行为。

# 在 mmpretrain/engine/optimizers/my_optim_constructor.py 中
from mmengine.optim import DefaultOptimWrapperConstructor
from mmpretrain.registry import OPTIM_WRAPPER_CONSTRUCTORS


@OPTIM_WRAPPER_CONSTRUCTORS.register_module()
class MyOptimWrapperConstructor:

    def __init__(self, optim_wrapper_cfg, paramwise_cfg=None):
        ...

    def __call__(self, model):
        ...

这是一个已实现的 OptimWrapperConstructor 具体例子。

接下来类似新增优化器教程来导入并使用新的优化器构造器。

修改 mmpretrain/engine/optimizers/__init__.py，将其导入至 mmpretrain.engine 包。

# 在 mmpretrain/engine/optimizers/__init__.py 中
...
from .my_optim_constructor import MyOptimWrapperConstructor

__all__ = [..., 'MyOptimWrapperConstructor']

在配置文件的 optim_wrapper.constructor 字段中使用 MyOptimWrapperConstructor 。

optim_wrapper = dict(
    constructor=dict(type='MyOptimWrapperConstructor'),
    optimizer=...,
    paramwise_cfg=...,
)