mirror of https://github.com/open-mmlab/mmcv.git
160 lines
6.1 KiB
Markdown
160 lines
6.1 KiB
Markdown
## 执行器
|
||
|
||
执行器模块负责模型训练过程调度，主要目的是让用户使用更少的代码以及灵活可配置方式开启训练。其具备如下核心特性:
|
||
|
||
- 支持以 `EpochBasedRunner` 和 `IterBasedRunner` 为单位的迭代模式以满足不同场景
|
||
- 支持定制工作流以满足训练过程中各状态自由切换，目前支持训练和验证两个工作流。工作流可以简单理解为一个完成的训练和验证迭代过程。
|
||
- 配合各类默认和自定义 Hook，对外提供了灵活扩展能力
|
||
|
||
### EpochBasedRunner
|
||
|
||
顾名思义，`EpochBasedRunner` 是指以 epoch 为周期的工作流，例如设置 workflow = \[('train', 2), ('val', 1)\] 表示循环迭代地训练 2 个 epoch，然后验证 1 个 epoch。MMDetection 目标检测框架默认采用的是 `EpochBasedRunner`。
|
||
|
||
其抽象逻辑如下所示：
|
||
|
||
```python
|
||
# 训练终止条件
|
||
while curr_epoch < max_epochs:
|
||
# 遍历用户设置的工作流，例如 workflow = [('train', 2)，('val', 1)]
|
||
for i, flow in enumerate(workflow):
|
||
# mode 是工作流函数，例如 train, epochs 是迭代次数
|
||
mode, epochs = flow
|
||
# 要么调用 self.train()，要么调用 self.val()
|
||
epoch_runner = getattr(self, mode)
|
||
# 运行对应工作流函数
|
||
for _ in range(epochs):
|
||
epoch_runner(data_loaders[i], **kwargs)
|
||
```
|
||
|
||
目前支持训练和验证两个工作流，以训练函数为例，其抽象逻辑是：
|
||
|
||
```python
|
||
# epoch_runner 目前可以是 train 或者 val
|
||
def train(self, data_loader, **kwargs):
|
||
# 遍历 dataset，共返回一个 epoch 的 batch 数据
|
||
for i, data_batch in enumerate(data_loader):
|
||
self.call_hook('before_train_iter')
|
||
# 验证时候 train_mode=False
|
||
self.run_iter(data_batch, train_mode=True, **kwargs)
|
||
self.call_hook('after_train_iter')
|
||
self.call_hook('after_train_epoch')
|
||
```
|
||
|
||
### IterBasedRunner
|
||
|
||
不同于 `EpochBasedRunner`，`IterBasedRunner` 是指以 iter 为周期的工作流，例如设置 workflow = \[('train', 2)， ('val', 1)\] 表示循环迭代的训练 2 个 iter，然后验证 1 个 iter，MMSegmentation 语义分割框架默认采用的是 `IterBasedRunner`。
|
||
|
||
其抽象逻辑如下所示：
|
||
|
||
```python
|
||
# 虽然是 iter 单位，但是某些场合需要 epoch 信息，由 IterLoader 提供
|
||
iter_loaders = [IterLoader(x) for x in data_loaders]
|
||
# 训练终止条件
|
||
while curr_iter < max_iters:
|
||
# 遍历用户设置的工作流，例如 workflow = [('train', 2)， ('val', 1)]
|
||
for i, flow in enumerate(workflow):
|
||
# mode 是工作流函数，例如 train, iters 是迭代次数
|
||
mode, iters = flow
|
||
# 要么调用 self.train()，要么调用 self.val()
|
||
iter_runner = getattr(self, mode)
|
||
# 运行对应工作流函数
|
||
for _ in range(iters):
|
||
iter_runner(iter_loaders[i], **kwargs)
|
||
```
|
||
|
||
目前支持训练和验证两个工作流，以验证函数为例，其抽象逻辑是：
|
||
|
||
```python
|
||
# iter_runner 目前可以是 train 或者 val
|
||
def val(self, data_loader, **kwargs):
|
||
# 获取 batch 数据，用于一次迭代
|
||
data_batch = next(data_loader)
|
||
self.call_hook('before_val_iter')
|
||
outputs = self.model.val_step(data_batch, self.optimizer, **kwargs)
|
||
self.outputs = outputs
|
||
self.call_hook('after_val_iter')
|
||
```
|
||
|
||
除了上述基础功能外，`EpochBasedRunner` 和 `IterBasedRunner` 还提供了 resume 、 save_checkpoint 和注册 hook 功能。
|
||
|
||
### 一个简单例子
|
||
|
||
以最常用的分类任务为例详细说明 `runner` 的使用方法。 开启任何一个训练任务，都需要包括如下步骤：
|
||
|
||
**(1) dataloader、model 和优化器等类初始化**
|
||
|
||
```python
|
||
# 模型类初始化
|
||
model=...
|
||
# 优化器类初始化，典型值 cfg.optimizer = dict(type='SGD', lr=0.1, momentum=0.9, weight_decay=0.0001)
|
||
optimizer = build_optimizer(model, cfg.optimizer)
|
||
# 工作流对应的 dataloader 初始化
|
||
data_loaders = [
|
||
build_dataloader(
|
||
ds,
|
||
cfg.data.samples_per_gpu,
|
||
cfg.data.workers_per_gpu,
|
||
...) for ds in dataset
|
||
]
|
||
```
|
||
|
||
**(2) runner 类初始化**
|
||
|
||
```python
|
||
runner = build_runner(
|
||
# cfg.runner 典型配置为
|
||
# runner = dict(type='EpochBasedRunner', max_epochs=200)
|
||
cfg.runner,
|
||
default_args=dict(
|
||
model=model,
|
||
batch_processor=None,
|
||
optimizer=optimizer,
|
||
logger=logger))
|
||
```
|
||
|
||
**(3) 注册默认训练所必须的 hook，和用户自定义 hook**
|
||
|
||
```python
|
||
# 注册定制必需的 hook
|
||
runner.register_training_hooks(
|
||
# lr相关配置，典型为
|
||
# lr_config = dict(policy='step', step=[100, 150])
|
||
cfg.lr_config,
|
||
# 优化相关配置，例如 grad_clip 等
|
||
optimizer_config,
|
||
# 权重保存相关配置，典型为
|
||
# checkpoint_config = dict(interval=1)，每个单位都保存权重
|
||
cfg.checkpoint_config,
|
||
# 日志相关配置
|
||
cfg.log_config,
|
||
...)
|
||
|
||
# 注册用户自定义 hook
|
||
# 例如想使用 ema 功能，则可以设置 custom_hooks=[dict(type='EMAHook')]
|
||
if cfg.get('custom_hooks', None):
|
||
custom_hooks = cfg.custom_hooks
|
||
for hook_cfg in cfg.custom_hooks:
|
||
hook_cfg = hook_cfg.copy()
|
||
priority = hook_cfg.pop('priority', 'NORMAL')
|
||
hook = build_from_cfg(hook_cfg, HOOKS)
|
||
runner.register_hook(hook, priority=priority)
|
||
```
|
||
|
||
然后可以进行 resume 或者 load_checkpoint 对权重进行加载。
|
||
|
||
**(4) 开启训练流**
|
||
|
||
```python
|
||
# workflow 典型为 workflow = [('train', 1)]
|
||
# 此时就真正开启了训练
|
||
runner.run(data_loaders, cfg.workflow)
|
||
```
|
||
|
||
关于 workflow 设置，以 `EpochBasedRunner` 为例，详情如下：
|
||
|
||
- 假设只想运行训练工作流，则可以设置 workflow = \[('train', 1)\]，表示只进行迭代训练
|
||
- 假设想运行训练和验证工作流，则可以设置 workflow = \[('train', 3), ('val', 1)\]，表示先训练 3 个 epoch ，然后切换到 val 工作流，运行 1 个 epoch，然后循环，直到训练 epoch 次数达到指定值
|
||
- 工作流设置还自由定制，例如你可以先验证再训练 workflow = \[('val', 1), ('train', 1)\]
|
||
|
||
上述代码都已经封装到了各个代码库的 train.py 中，用户只需要设置相应的配置即可，上述流程会自动运行。
|