mmcv文件夹
runner
执行器模块负责模型训练过程调度,主要目的是让用户使用更少的代码以及灵活可配置方式开启训练。其具备如下核心特性:
支持以 EpochBasedRunner 和 IterBasedRunner 为单位的迭代模式以满足不同场景
支持定制工作流以满足训练过程中各状态自由切换,目前支持训练和验证两个工作流。工作流可以简单理解为一个完成的训练和验证迭代过程。
配合各类默认和自定义 Hook,对外提供了灵活扩展能力
EpochBasedRunner(mmcv/runner/epoch_based_runner.py)
顾名思义,EpochBasedRunner 是指以 epoch 为周期的工作流,例如设置 workflow = [(‘train’, 2), (‘val’, 1)] 表示循环迭代地训练 2 个 epoch,然后验证 1 个 epoch。MMDetection 目标检测框架默认采用的是 EpochBasedRunner。
其抽象逻辑如下所示:
# 训练终止条件
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)
目前支持训练和验证两个工作流,以训练函数为例,其抽象逻辑是:
# 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(mmcv/runner/iter_based_runner.py)
不同于 EpochBasedRunner,IterBasedRunner 是指以 iter 为周期的工作流,例如设置 workflow = [(‘train’, 2), (‘val’, 1)] 表示循环迭代的训练 2 个 iter,然后验证 1 个 iter,MMSegmentation 语义分割框架默认采用的是 IterBasedRunner。
其抽象逻辑如下所示:
# 虽然是 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)
# 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 和优化器等类初始化
# 模型类初始化
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
]