customize_runtime.md 13.7 KB
Newer Older
1
# 自定义运行时配置
2
3
4

## 自定义优化器设置

5
优化器相关的配置是由 `optim_wrapper` 管理的,其通常有三个字段:`optimizer``paramwise_cfg``clip_grad`。更多细节请参考 [OptimWrapper](https://mmengine.readthedocs.io/zh_CN/latest/tutorials/optim_wrapper.html)。如下所示,使用 `AdamW` 作为`优化器`,骨干网络的学习率降低 10 倍,并添加了梯度裁剪。
6
7
8
9

```python
optim_wrapper = dict(
    type='OptimWrapper',
10
    # 优化器
11
12
13
14
15
16
17
    optimizer=dict(
        type='AdamW',
        lr=0.0001,
        weight_decay=0.05,
        eps=1e-8,
        betas=(0.9, 0.999)),

18
    # 参数级学习率及权重衰减系数设置
19
20
21
22
23
24
    paramwise_cfg=dict(
        custom_keys={
            'backbone': dict(lr_mult=0.1, decay_mult=1.0),
        },
        norm_decay_mult=0.0),

25
    # 梯度裁剪
26
27
28
    clip_grad=dict(max_norm=0.01, norm_type=2))
```

29
30
### 自定义 PyTorch 支持的优化器

31
我们已经支持使用所有 PyTorch 实现的优化器,且唯一需要修改的地方就是改变配置文件中的 `optim_wrapper` 字段中的 `optimizer` 字段。例如,如果您想使用 `ADAM`(注意到这样可能会使性能大幅下降),您可以这样修改:
32
33

```python
34
35
36
optim_wrapper = dict(
    type='OptimWrapper',
    optimizer=dict(type='Adam', lr=0.0003, weight_decay=0.0001))
37
38
```

39
为了修改模型的学习率,用户只需要修改 `optimizer` 中的 `lr` 字段。用户可以根据 PyTorch 的 [API 文档](https://pytorch.org/docs/stable/optim.html?highlight=optim#module-torch.optim)直接设置参数。
40
41
42
43
44
45
46

### 自定义并实现优化器

#### 1. 定义新的优化器

一个自定义优化器可以按照如下过程定义:

47
假设您想要添加一个叫 `MyOptimizer` 的,拥有参数 `a``b``c` 的优化器,您需要创建一个叫做 `mmdet3d/engine/optimizers` 的目录。接下来,应该在目录下某个文件中实现新的优化器,比如 `mmdet3d/engine/optimizers/my_optimizer.py`
48
49

```python
50
from mmdet3d.registry import OPTIMIZERS
51
52
53
54
55
56
from torch.optim import Optimizer


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

57
    def __init__(self, a, b, c):
58
59
60
61
62

```

#### 2. 将优化器添加到注册器

63
为了找到上述定义的优化器模块,该模块首先需要被引入主命名空间。有两种实现方法:
64

65
- 修改 `mmdet3d/engine/optimizers/__init__.py` 导入该模块。
66

67
  新定义的模块应该在 `mmdet3d/engine/optimizers/__init__.py` 中被导入,从而被找到并且被添加到注册器中:
68

69
70
71
  ```python
  from .my_optimizer import MyOptimizer
  ```
72

73
- 在配置中使用 `custom_imports` 来人工导入新优化器:
74

75
76
77
  ```python
  custom_imports = dict(imports=['mmdet3d.engine.optimizers.my_optimizer'], allow_failed_imports=False)
  ```
78

79
模块 `mmdet3d.engine.optimizers.my_optimizer` 会在程序开始被导入,且 `MyOptimizer` 类在那时会自动被注册。注意到应该只有包含 `MyOptimizer` 类的包被导入。`mmdet3d.engine.optimizers.my_optimizer.MyOptimizer`**不能**被直接导入。
80

81
事实上,用户可以在这种导入的方法中使用完全不同的文件目录结构,只要保证根目录能在 `PYTHONPATH` 中被定位。
82
83
84

#### 3. 在配置文件中指定优化器

85
接下来您可以在配置文件的 `optimizer` 字段中使用 `MyOptimizer`。在配置文件中,优化器在 `optimizer` 字段中以如下方式定义:
86
87

```python
88
89
90
optim_wrapper = dict(
    type='OptimWrapper',
    optimizer=dict(type='SGD', lr=0.02, momentum=0.9, weight_decay=0.0001))
91
92
93
94
95
```

为了使用您自己的优化器,该字段可以改为:

```python
96
97
98
optim_wrapper = dict(
    type='OptimWrapper',
    optimizer=dict(type='MyOptimizer', a=a_value, b=b_value, c=c_value))
99
100
```

101
### 自定义优化器封装构造器
102

103
部分模型可能会拥有一些参数专属的优化器设置,比如 BatchNorm 层的权重衰减 (weight decay)。用户可以通过自定义优化器封装构造器来对那些细粒度的参数进行调优。
104
105

```python
106
from mmengine.optim import DefaultOptiWrapperConstructor
107

108
from mmdet3d.registry import OPTIM_WRAPPER_CONSTRUCTORS
109
110
111
from .my_optimizer import MyOptimizer


112
113
@OPTIM_WRAPPER_CONSTRUCTORS.register_module()
class MyOptimizerWrapperConstructor(DefaultOptimWrapperConstructor):
114

115
116
117
    def __init__(self,
                 optim_wrapper_cfg: dict,
                 paramwise_cfg: Optional[dict] = None):
118

119
    def __call__(self, model: nn.Module) -> OptimWrapper:
120

121
        return optim_wrapper
122
123
124

```

125
默认优化器封装构造器在[这里](https://github.com/open-mmlab/mmengine/blob/main/mmengine/optim/optimizer/default_constructor.py#L18)实现。这部分代码也可以用作新优化器封装构造器的模版。
126
127
128

### 额外的设置

129
没有在优化器部分实现的技巧应该通过优化器封装构造器或者钩子来实现(比如逐参数的学习率设置)。我们列举了一些常用的可以稳定训练过程或者加速训练的设置。我们欢迎提供更多类似设置的 PR 和 issue。
130
131
132

- __使用梯度裁剪 (gradient clip) 来稳定训练过程__:

133
  一些模型依赖梯度裁剪技术来裁剪训练中的梯度,以稳定训练过程。举例如下:
134

135
  ```python
136
137
  optim_wrapper = dict(
      _delete_=True, clip_grad=dict(max_norm=35, norm_type=2))
138
  ```
139

140
  如果您的配置继承了一个已经设置了 `optim_wrapper` 的基础配置,那么您可能需要 `_delete_=True` 字段来覆盖基础配置中无用的设置。更多细节请参考[说明文档](https://mmdetection3d.readthedocs.io/en/latest/tutorials/config.html)
141

142
- __使用动量调度器 (momentum scheduler) 来加速模型收敛__:
143

144
  我们支持用动量调度器来根据学习率更改模型的动量,这样可以使模型更快地收敛。动量调度器通常和学习率调度器一起使用,例如,如下配置文件在 [3D 检测](https://github.com/open-mmlab/mmdetection3d/blob/dev-1.x/configs/_base_/schedules/cyclic_20e.py)中被用于加速模型收敛。更多细节请参考 [CosineAnnealingLR](https://github.com/open-mmlab/mmengine/blob/main/mmengine/optim/scheduler/lr_scheduler.py#L43)[CosineAnnealingMomentum](https://github.com/open-mmlab/mmengine/blob/main/mmengine/optim/scheduler/momentum_scheduler.py#L71) 的实现方法。
145
146

  ```python
147
  param_scheduler = [
148
149
150
      # 学习率调度器
      # 在前 8 个 epoch,学习率从 0 升到 lr * 10
      # 在接下来 12 个 epoch,学习率从 lr * 10 降到 lr * 1e-4
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
      dict(
          type='CosineAnnealingLR',
          T_max=8,
          eta_min=lr * 10,
          begin=0,
          end=8,
          by_epoch=True,
          convert_to_iter_based=True),
      dict(
          type='CosineAnnealingLR',
          T_max=12,
          eta_min=lr * 1e-4,
          begin=8,
          end=20,
          by_epoch=True,
          convert_to_iter_based=True),
167
168
169
      # 动量调度器
      # 在前 8 个 epoch,动量从 0 升到 0.85 / 0.95
      # 在接下来 12 个 epoch,动量从 0.85 / 0.95 升到 1
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
      dict(
          type='CosineAnnealingMomentum',
          T_max=8,
          eta_min=0.85 / 0.95,
          begin=0,
          end=8,
          by_epoch=True,
          convert_to_iter_based=True),
      dict(
          type='CosineAnnealingMomentum',
          T_max=12,
          eta_min=1,
          begin=8,
          end=20,
          by_epoch=True,
          convert_to_iter_based=True)
  ]
187
  ```
188

189
## 自定义训练调度
190

191
默认情况下我们使用阶梯式学习率衰减的 1 倍训练调度。这会调用 MMEngine 中的 [`MultiStepLR`](https://github.com/open-mmlab/mmengine/blob/main/mmengine/optim/scheduler/lr_scheduler.py#L139)。我们在[这里](https://github.com/open-mmlab/mmengine/blob/main/mmengine/optim/scheduler/lr_scheduler.py)支持了很多其他学习率调度,比如`余弦退火``多项式衰减`调度。下面是一些样例:
192

193
- 多项式衰减调度:
194

195
  ```python
196
197
198
199
200
201
202
203
  param_scheduler = [
      dict(
          type='PolyLR',
          power=0.9,
          eta_min=1e-4,
          begin=0,
          end=8,
          by_epoch=True)]
204
  ```
205

206
- 余弦退火调度:
207

208
  ```python
209
210
211
212
213
214
215
216
  param_scheduler = [
      dict(
          type='CosineAnnealingLR',
          T_max=8,
          eta_min=lr * 1e-5,
          begin=0,
          end=8,
          by_epoch=True)]
217
  ```
218

219
## 自定义训练循环控制器
220

221
默认情况下,我们在 `train_cfg` 中使用 `EpochBasedTrainLoop`,并在每一个训练 epoch 完成后进行一次验证,如下所示:
222
223

```python
224
train_cfg = dict(type='EpochBasedTrainLoop', max_epochs=12, val_begin=1, val_interval=1)
225
226
```

227
事实上,[`IterBasedTrainLoop`](https://github.com/open-mmlab/mmengine/blob/main/mmengine/runner/loops.py#L183%5D)[`EpochBasedTrainLoop`](https://github.com/open-mmlab/mmengine/blob/main/mmengine/runner/loops.py#L18) 都支持动态间隔验证,如下所示:
228
229

```python
230
231
232
# 在第 365001 次迭代之前,我们每隔 5000 次迭代验证一次
# 在第 365000 次迭代之后,我们每隔 368750 次迭代验证一次
# 这意味着我们在训练结束后进行验证。
233
234
235
236
237
238
239
240
241

interval = 5000
max_iters = 368750
dynamic_intervals = [(max_iters // interval * interval + 1, max_iters)]
train_cfg = dict(
    type='IterBasedTrainLoop',
    max_iters=max_iters,
    val_interval=interval,
    dynamic_intervals=dynamic_intervals)
242
243
244
245
246
247
248
249
```

## 自定义钩子

### 自定义并实现钩子

#### 1. 实现一个新钩子

250
MMEngine 提供了一些实用的[钩子](https://github.com/open-mmlab/mmengine/blob/main/docs/en/tutorials/hook.md),但有些场合用户可能需要实现一个新的钩子。在 v1.1.0rc0 之后,MMDetection3D 在训练时支持基于 MMEngine 自定义钩子。因此用户可以直接在 mmdet3d 或者基于 mmdet3d 的代码库中实现钩子并通过更改训练配置来使用钩子。这里我们给出一个在 mmdet3d 中创建并使用新钩子的例子。
251
252

```python
253
254
255
from mmengine.hooks import Hook

from mmdet3d.registry import HOOKS
256
257
258
259
260
261
262


@HOOKS.register_module()
class MyHook(Hook):

    def __init__(self, a, b):

263
    def before_run(self, runner) -> None:
264

265
    def after_run(self, runner) -> None:
266

267
    def before_train(self, runner) -> None:
268

269
    def after_train(self, runner) -> None:
270

271
    def before_train_epoch(self, runner) -> None:
272

273
274
275
276
277
278
279
280
281
282
283
284
    def after_train_epoch(self, runner) -> None:

    def before_train_iter(self,
                          runner,
                          batch_idx: int,
                          data_batch: DATA_BATCH = None) -> None:

    def after_train_iter(self,
                         runner,
                         batch_idx: int,
                         data_batch: DATA_BATCH = None,
                         outputs: Optional[dict] = None) -> None:
285
286
```

287
用户需要根据钩子的功能指定钩子在每个训练阶段时的行为,具体包括如下阶段:`before_run``after_run``before_train``after_train``before_train_epoch``after_train_epoch``before_train_iter`,和 `after_train_iter`。有更多的位点可以插入钩子,详情可参考 [base hook class](https://github.com/open-mmlab/mmengine/blob/main/mmengine/hooks/hook.py#L9)
288
289
290

#### 2. 注册新钩子

291
接下来我们需要导入 `MyHook`。假设新钩子位于文件 `mmdet3d/engine/hooks/my_hook.py` 中,有两种实现方法:
292

293
- 修改 `mmdet3d/engine/hooks/__init__.py` 导入该模块:
294

295
  新定义的模块应该在 `mmdet3d/engine/hooks/__init__.py` 中被导入,从而被找到并且被添加到注册器中:
296

297
298
299
  ```python
  from .my_hook import MyHook
  ```
300

301
- 在配置中使用 `custom_imports` 来人为地导入新钩子:
302

303
304
305
  ```python
  custom_imports = dict(imports=['mmdet3d.engine.hooks.my_hook'], allow_failed_imports=False)
  ```
306
307
308
309
310
311
312
313
314

#### 3. 更改配置文件

```python
custom_hooks = [
    dict(type='MyHook', a=a_value, b=b_value)
]
```

315
您可以将字段 `priority` 设置为 `'NORMAL'` 或者 `'HIGHEST'` 来设置钩子的优先级,如下所示:
316
317
318
319
320
321
322

```python
custom_hooks = [
    dict(type='MyHook', a=a_value, b=b_value, priority='NORMAL')
]
```

323
默认情况下,注册阶段钩子的优先级为 `NORMAL`
324

325
### 使用 MMEngine 中实现的钩子
326

327
如果钩子已经在 MMEngine 中被实现了,您可以直接通过更改配置文件来使用该钩子:
328
329
330

### 更改默认的运行时钩子

331
有一些常用的钩子通过 `default_hooks` 注册,它们是:
332

333
334
335
336
337
- `IterTimerHook`:该钩子用来记录加载数据的时间 'data_time' 和模型训练一步的时间 'time' 。
- `LoggerHook`:该钩子用来从`执行器(Runner)`的不同组件收集日志并将其写入终端,Json 文件,tensorboard 和 wandb 等。
- `ParamSchedulerHook`:该钩子用来更新优化器中的一些超参数,例如学习率和动量。
- `CheckpointHook`:该钩子用来定期地保存检查点。
- `DistSamplerSeedHook`:该钩子用来设置采样和批采样的种子。
338

339
`IterTimerHook``ParamSchedulerHook``DistSamplerSeedHook` 都很简单,通常不需要修改,因此此处我们将介绍如何使用 `LoggerHook``CheckpointHook``Det3DVisualizationHook`
340

341
#### CheckpointHook
342

343
除了定期地保存检查点,[`CheckpointHook`](https://github.com/open-mmlab/mmengine/blob/main/mmengine/hooks/checkpoint_hook.py#L19) 提供了其它的可选项例如 `max_keep_ckpts``save_optimizer` 等。用户可以设置 `max_keep_ckpts` 只保存少量的检查点或者通过 `save_optimizer` 决定是否保存优化器的状态。参数的更多细节参考[此处](https://github.com/open-mmlab/mmengine/blob/main/mmengine/hooks/checkpoint_hook.py#L19)
344
345

```python
346
347
348
349
350
351
default_hooks = dict(
    checkpoint=dict(
        type='CheckpointHook',
        interval=1,
        max_keep_ckpts=3,
        save_optimizer=True))
352
353
```

354
#### LoggerHook
355

356
`LoggerHook` 允许设置日志记录间隔。详细介绍可参考[文档](https://github.com/open-mmlab/mmengine/blob/main/mmengine/hooks/logger_hook.py#L18)
357
358

```python
359
default_hooks = dict(logger=dict(type='LoggerHook', interval=50))
360
```