371 lines
15 KiB
Markdown
371 lines
15 KiB
Markdown
# 自定义训练优化策略
|
||
|
||
在我们的算法库中，已经提供了通用数据集（如ImageNet，CIFAR）的[默认训练策略配置](https://github.com/open-mmlab/mmclassification/blob/master/configs/_base_/schedules)。如果想要在这些数据集上继续提升模型性能，或者在不同数据集和方法上进行新的尝试，我们通常需要修改这些默认的策略。
|
||
|
||
在本教程中，我们将介绍如何在运行自定义训练时，通过修改配置文件进行构造优化器、参数化精细配置、梯度裁剪、梯度累计以及定制动量调整策略等。同时也会通过模板简单介绍如何自定义开发优化器和构造器。
|
||
|
||
## 配置训练优化策略
|
||
|
||
我们通过 `optim_wrapper` 来配置主要的优化策略，包括优化器的选择，混合精度训练的选择，参数化精细配置，梯度裁剪以及梯度累计。接下来将分别介绍这些内容。
|
||
|
||
### 构造 PyTorch 内置优化器
|
||
|
||
MMClassification 支持 PyTorch 实现的所有优化器，仅需在配置文件中，指定优化器封装需要的 `optimizer` 字段。
|
||
|
||
如果要使用 [`SGD`](torch.optim.SGD)，则修改如下。这里要注意所有优化相关的配置都需要封装在 `optim_wrapper` 配置里。
|
||
|
||
```python
|
||
optim_wrapper = dict(
|
||
type='OptimWrapper',
|
||
optimizer=dict(type='SGD', lr=0.0003, weight_decay=0.0001)
|
||
)
|
||
```
|
||
|
||
```{note}
|
||
配置文件中的 'type' 不是构造时的参数，而是 PyTorch 内置优化器的类名。
|
||
更多优化器选择可以参考{external+torch:ref}`PyTorch 支持的优化器列表<optim:algorithms>`。
|
||
```
|
||
|
||
要修改模型的学习率，只需要在优化器的配置中修改 `lr` 即可。
|
||
要配置其他参数，可直接根据 [PyTorch API 文档](torch.optim) 进行。
|
||
|
||
例如，如果想使用 [`Adam`](torch.optim.Adam) 并设置参数为 `torch.optim.Adam(params, lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0, amsgrad=False)`。
|
||
则需要进行如下修改：
|
||
|
||
```python
|
||
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),
|
||
)
|
||
```
|
||
|
||
````{note}
|
||
考虑到对于单精度训练来说，优化器封装的默认类型就是 `OptimWrapper`，我们在这里可以直接省略，因此配置文件可以进一步简化为：
|
||
|
||
```python
|
||
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`。
|
||
|
||
```python
|
||
optim_wrapper = dict(type='AmpOptimWrapper', optimizer=...)
|
||
```
|
||
|
||
另外，为了方便，我们同时在启动训练脚本 `tools/train.py` 中提供了 `--amp` 参数作为开启混合精度训练的开关，更多细节可以参考[训练教程](../user_guides/train.md)。
|
||
|
||
### 参数化精细配置
|
||
|
||
在一些模型中，不同的优化策略需要适应特定的参数，例如不在 BatchNorm 层使用权重衰减，或者在不同层使用不同的学习率等等。
|
||
我们需要用到 `optim_wrapper` 中的 `paramwise_cfg` 参数来进行精细化配置。
|
||
|
||
- **为不同类型的参数设置超参乘子**
|
||
|
||
例如，我们可以在 `paramwise_cfg` 配置中设置 `norm_decay_mult=0.` 来改变归一化层权重和偏移的衰减为0。
|
||
|
||
```python
|
||
optim_wrapper = dict(
|
||
optimizer=dict(type='SGD', lr=0.8, weight_decay=1e-4),
|
||
paramwise_cfg=dict(norm_decay_mult=0.))
|
||
```
|
||
|
||
支持更多类型的参数配置，参考以下列表：
|
||
|
||
- `lr_mult`：所有参数的学习率系数
|
||
- `decay_mult`：所有参数的衰减系数
|
||
- `bias_lr_mult`：偏置的学习率系数（不包括正则化层的偏置以及可变形卷积的 offset），默认值为 1
|
||
- `bias_decay_mult`：偏置的权值衰减系数（不包括正则化层的偏置以及可变形卷积的 offset），默认值为 1
|
||
- `norm_decay_mult`：正则化层权重和偏置的权值衰减系数，默认值为 1
|
||
- `dwconv_decay_mult`：Depth-wise 卷积的权值衰减系数，默认值为 1
|
||
- `bypass_duplicate`：是否跳过重复的参数，默认为 `False`
|
||
- `dcn_offset_lr_mult`：可变形卷积（Deformable Convolution）的学习率系数，默认值为 1
|
||
|
||
- **为特定参数设置超参乘子**
|
||
|
||
MMClassification 通过 `paramwise_cfg` 的 `custom_keys` 参数来配置特定参数的超参乘子。
|
||
|
||
例如，我们可以通过以下配置来设置所有 `backbone.layer0` 层的学习率和权重衰减为0， `backbone` 的其余层和优化器保持一致，另外 `head` 层的学习率为0.001.
|
||
|
||
```python
|
||
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)
|
||
}))
|
||
```
|
||
|
||
### 梯度裁剪
|
||
|
||
在训练过程中，损失函数可能接近于一些异常陡峭的区域，从而导致梯度爆炸。而梯度裁剪可以帮助稳定训练过程，更多介绍可以参见[该页面](https://paperswithcode.com/method/gradient-clipping)。
|
||
|
||
目前我们支持在 `optim_wrapper` 字段中添加 `clip_grad` 参数来进行梯度裁剪，更详细的参数可参考 [PyTorch 文档](torch.nn.utils.clip_grad_norm_)。
|
||
|
||
用例如下：
|
||
|
||
```python
|
||
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` 参数来进行梯度累计。
|
||
|
||
用例如下：
|
||
|
||
```python
|
||
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，也即：
|
||
|
||
```python
|
||
train_dataloader = dict(batch_size=256)
|
||
optim_wrapper = dict(
|
||
optimizer=dict(type='SGD', lr=0.01, weight_decay=0.0001))
|
||
```
|
||
|
||
## 配置参数优化策略
|
||
|
||
在训练过程中，优化参数例如学习率、动量，通常不会是固定不变，而是随着训练进程的变化而调整。PyTorch 支持一些学习率调整的调度器，但是不足以完成复杂的策略。在MMClassification中，我们提供 `param_scheduler` 来更好地控制不同优化参数的策略。
|
||
|
||
### 配置学习率调整策略
|
||
|
||
深度学习研究中，广泛应用学习率衰减来提高网络的性能。我们支持大多数 PyTorch 学习率调度器， 其中包括 `ExponentialLR`, `LinearLR`, `StepLR`, `MultiStepLR` 等等。
|
||
|
||
- **单个学习率策略**
|
||
|
||
多数情况下，我们使用单一学习率策略，这里 `param_scheduler` 会是一个字典。比如在默认的 ResNet 网络训练中，我们使用阶梯式的学习率衰减策略 [`MultiStepLR`](mmengine.optim.MultiStepLR)，配置文件为：
|
||
|
||
```python
|
||
param_scheduler = dict(
|
||
type='MultiStepLR',
|
||
by_epoch=True,
|
||
milestones=[100, 150],
|
||
gamma=0.1)
|
||
```
|
||
|
||
或者我们想使用 [`CosineAnnealingLR`](mmengine.optim.CosineAnnealingLR) 来进行学习率衰减：
|
||
|
||
```python
|
||
param_scheduler = dict(
|
||
type='CosineAnnealingLR',
|
||
by_epoch=True,
|
||
T_max=num_epochs)
|
||
```
|
||
|
||
- **多个学习率策略**
|
||
|
||
然而在一些其他情况下，为了提高模型的精度，通常会使用多种学习率策略。例如，在训练的早期阶段，网络容易不稳定，而学习率的预热就是为了减少这种不稳定性。
|
||
|
||
整个学习过程中，学习率将会通过预热从一个很小的值逐步提高到预定值，再会通过其他的策略进一步调整。
|
||
|
||
在 MMClassification 中，我们同样使用 `param_scheduler` ，将多种学习策略写成列表就可以完成上述预热策略的组合。
|
||
|
||
例如：
|
||
|
||
1. 在前50次迭代中逐**迭代次数**地**线性**预热
|
||
|
||
```python
|
||
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)
|
||
]
|
||
```
|
||
|
||
2. 在前10轮迭代中逐**迭代次数**地**线性**预热
|
||
|
||
```python
|
||
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`](torch.optim.lr_scheduler.ChainedScheduler) 一致）。
|
||
|
||
```{tip}
|
||
为了避免学习率曲线与预期不符， 配置完成后，可以使用 MMClassification 提供的 [学习率可视化工具](../useful_tools/scheduler_visualization.md) 画出对应学习率调整曲线。
|
||
```
|
||
|
||
### 配置动量调整策略
|
||
|
||
MMClassification 支持动量调度器根据学习率修改优化器的动量，从而使损失函数收敛更快。用法和学习率调度器一致。
|
||
|
||
我们支持的动量策略和详细的使用细节可以参考[这里](https://github.com/open-mmlab/mmengine/blob/main/mmengine/optim/scheduler/momentum_scheduler.py)。我们只将调度器中的 `LR` 替换为了 `Momentum`，动量策略可以直接追加 `param_scheduler` 列表中。
|
||
|
||
这里是一个用例：
|
||
|
||
```python
|
||
param_scheduler = [
|
||
# 学习率策略
|
||
dict(type='LinearLR', ...),
|
||
# 动量策略
|
||
dict(type='LinearMomentum',
|
||
start_factor=0.001,
|
||
by_epoch=False,
|
||
begin=0,
|
||
end=1000)
|
||
]
|
||
```
|
||
|
||
## 新增优化器或者优化器构造器
|
||
|
||
```{note}
|
||
本部分将修改 MMClassification 源码或者向 MMClassification 框架添加代码，初学者可跳过。
|
||
```
|
||
|
||
### 新增优化器
|
||
|
||
在学术研究和工业实践中，可能需要使用 MMClassification 未实现的优化方法，可以通过以下方法添加。
|
||
|
||
#### 1. 定义一个新的优化器
|
||
|
||
一个自定义的优化器可根据如下规则进行定制：
|
||
|
||
假设我们想添加一个名为 `MyOptimzer` 的优化器，其拥有参数 `a`, `b` 和 `c`。
|
||
可以创建一个名为 `mmpretrain/engine/optimizer` 的文件夹，并在目录下的一个文件，如 `mmpretrain/engine/optimizer/my_optimizer.py` 中实现该自定义优化器：
|
||
|
||
```python
|
||
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):
|
||
...
|
||
```
|
||
|
||
#### 2. 注册优化器
|
||
|
||
要注册上面定义的上述模块，首先需要将此模块导入到主命名空间中。有两种方法可以实现它。
|
||
|
||
- 修改 `mmpretrain/engine/optimizers/__init__.py`，将其导入至 `mmpretrain.engine` 包。
|
||
|
||
```python
|
||
# 在 mmpretrain/engine/optimizers/__init__.py 中
|
||
...
|
||
from .my_optimizer import MyOptimizer # MyOptimizer 是我们自定义的优化器的名字
|
||
|
||
__all__ = [..., 'MyOptimizer']
|
||
```
|
||
|
||
在运行过程中，我们会自动导入 `mmpretrain.engine` 包并同时注册 `MyOptimizer`。
|
||
|
||
- 在配置中使用 `custom_imports` 手动导入
|
||
|
||
```python
|
||
custom_imports = dict(
|
||
imports=['mmpretrain.engine.optimizers.my_optimizer'],
|
||
allow_failed_imports=False,
|
||
)
|
||
```
|
||
|
||
`mmpretrain.engine.optimizers.my_optimizer` 模块将会在程序开始阶段被导入，`MyOptimizer` 类会随之自动被注册。
|
||
注意，这里只需要导入包含 `MyOptmizer` 类的包。如果填写`mmpretrain.engine.optimizers.my_optimizer.MyOptimizer` 则 **不会** 被直接导入。
|
||
|
||
#### 3. 在配置文件中指定优化器
|
||
|
||
之后，用户便可在配置文件的 `optim_wrapper.optimizer` 域中使用 `MyOptimizer`：
|
||
|
||
```python
|
||
optim_wrapper = dict(
|
||
optimizer=dict(type='MyOptimizer', a=a_value, b=b_value, c=c_value))
|
||
```
|
||
|
||
### 新增优化器构造器
|
||
|
||
某些模型可能具有一些特定于参数的设置以进行优化，例如 为所有 BatchNorm 层设置不同的权重衰减。
|
||
|
||
尽管我们已经可以使用 [`optim_wrapper.paramwise_cfg` 字段](#参数化精细配置)来配置特定参数的优化设置，但可能仍然无法覆盖你的需求。
|
||
|
||
当然你可以在此基础上进行修改。我们默认使用 [`DefaultOptimWrapperConstructor`](mmengine.optim.DefaultOptimWrapperConstructor) 来构造优化器。在构造过程中，通过 `paramwise_cfg` 来精细化配置不同设置。这个默认构造器可以作为新优化器构造器实现的模板。
|
||
|
||
我们可以新增一个优化器构造器来覆盖这些行为。
|
||
|
||
```python
|
||
# 在 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):
|
||
...
|
||
```
|
||
|
||
接下来类似 [新增优化器教程](#新增优化器) 来导入并使用新的优化器构造器。
|
||
|
||
1. 修改 `mmpretrain/engine/optimizers/__init__.py`，将其导入至 `mmpretrain.engine` 包。
|
||
|
||
```python
|
||
# 在 mmpretrain/engine/optimizers/__init__.py 中
|
||
...
|
||
from .my_optim_constructor import MyOptimWrapperConstructor
|
||
|
||
__all__ = [..., 'MyOptimWrapperConstructor']
|
||
```
|
||
|
||
2. 在配置文件的 `optim_wrapper.constructor` 字段中使用 `MyOptimWrapperConstructor` 。
|
||
|
||
```python
|
||
optim_wrapper = dict(
|
||
constructor=dict(type='MyOptimWrapperConstructor'),
|
||
optimizer=...,
|
||
paramwise_cfg=...,
|
||
)
|
||
```
|