mmdetection详解指北 (四)

发表于 | 分类于

训练流程

训练流程的包装过程大致如下:tools/train.py->apis/train.py->mmcv/runner.py->mmcv/hook.py(后面是分散的), 其中 runner 维护了数据信息,优化器, 日志系统, 训练 loop 中的各节点信息, 模型保存, 学习率等. 另外补充一点, 以上包装过程, 在 mmdet 中无处不在, 包括 mmcv 的代码也是对日常频繁使用的函数进行了统一封装。

训练逻辑

image-20200917203908378

图见2, 注意它的四个层级. 代码上, 主要查看 apis/train.py, mmcv 中的runner 相关文件. 核心围绕 Runner,Hook 两个类. Runner 将模型, 批处理函数 batch_pro cessor, 优化器作为基本属性, 训练过程中与训练状态, 各节点相关的信息被记录在mode,_hooks,_epoch,_iter,_inner_iter,_max_epochs, _max_iters中，这些信息维护了训练过程中插入不同 hook 的操作方式. 理清训练流程只需看 Runner 的成员函数 run. 在 run 里会根据 mode 按配置中 workflow 的 epoch 循环调用 train 和 val 函数, 跑完所有的 epoch. 比如train:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
class EpochBasedRunner(BaseRunner):
    """Epoch-based Runner.

    This runner train models epoch by epoch.
    """

    def train(self, data_loader, **kwargs):
        self.model.train()
        self.mode = 'train'
        self.data_loader = data_loader
        self._max_iters = self._max_epochs * len(self.data_loader) # 最大迭代次数
        self.call_hook('before_train_epoch')# 根据名字获取hook对象函数
        time.sleep(2)  # Prevent possible deadlock during epoch transition
        for i, data_batch in enumerate(self.data_loader):
            self._inner_iter = i # 记录当前训练迭代次数
            self.call_hook('before_train_iter') #一个batch 前向操作开始
            if self.batch_processor is None:
                outputs = self.model.train_step(data_batch, self.optimizer,
                                                **kwargs)
            else:
                outputs = self.batch_processor(
                    self.model, data_batch, train_mode=True, **kwargs) 
            if not isinstance(outputs, dict):
                raise TypeError('"batch_processor()" or "model.train_step()"'
                                ' must return a dict')
            if 'log_vars' in outputs:
                self.log_buffer.update(outputs['log_vars'],
                                       outputs['num_samples'])
            self.outputs = outputs
            self.call_hook('after_train_iter')#一个batch 前向操作结束
            self._iter += 1 # 方便resume，知道从哪次开始优化

        self.call_hook('after_train_epoch') #一个epoch结束
        self._epoch += 1 #记录训练epoch状态，方便resume

上面需要说明的是自定义 hook 类, 自定义 hook 类需继承 mmcv 的Hook 类, 其默认了 6+8+4 个成员函数, 也即2所示的 6 个层级节点, 外加 2*4 个区分 train 和 val 的节点记录函数, 以及 4 个边界检查函数. 从train.py 中容易看出, 在训练之前, 已经将需要的 hook 函数注册到 Runner的 self._hook 中了, 包括从配置文件解析的优化器, 学习率调整函数, 模型保存, 一个 batch 的时间记录等 (注册 hook 算子在 self._hook 中按优先级升序排列). 这里的 call_hook 函数定义如下:

1
2
3
4
5
6
7
8
9
10
11
12
class BaseRunner(metaclass=ABCMeta):
    ...
    def call_hook(self, fn_name):
        """Call all hooks.

        Args:
            fn_name (str): The function name in each hook to be called, such as
                "before_train_epoch".
        """
        for hook in self._hooks:
            getattr(hook, fn_name)(self)
    ...

容易看出, 在训练的不同节点, 将从注册列表中调用实现了该节点函数的类成员函数. 比如

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
@HOOKS.register_module()
class OptimizerHook(Hook):

    def __init__(self, grad_clip=None):
        self.grad_clip = grad_clip

    def clip_grads(self, params):
        params = list(
            filter(lambda p: p.requires_grad and p.grad is not None, params))
        if len(params) > 0:
            return clip_grad.clip_grad_norm_(params, **self.grad_clip)

    def after_train_iter(self, runner):
        runner.optimizer.zero_grad()
        runner.outputs['loss'].backward()
        if self.grad_clip is not None:
            grad_norm = self.clip_grads(runner.model.parameters())
            if grad_norm is not None:
                # Add grad norm to the logger
                runner.log_buffer.update({'grad_norm': float(grad_norm)},
                                         runner.outputs['num_samples'])
        runner.optimizer.step()

将在每个 train_iter 后实现反向传播和参数更新。学习率优化相对复杂一点, 其基类 LrUpdaterHook, 实现了 before_run,before_train_epoch, before_train_iter 三个 hook 函数, 意义自明. . 这里选一个余弦式变化, 稍作说明:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
def annealing_cos(start, end, factor, weight=1):
    """Calculate annealing cos learning rate.

    Cosine anneal from `weight * start + (1 - weight) * end` to `end` as
    percentage goes from 0.0 to 1.0.

    Args:
        start (float): The starting learning rate of the cosine annealing.
        end (float): The ending learing rate of the cosine annealing.
        factor (float): The coefficient of `pi` when calculating the current
            percentage. Range from 0.0 to 1.0.
        weight (float, optional): The combination factor of `start` and `end`
            when calculating the actual starting learning rate. Default to 1.
    """
    cos_out = cos(pi * factor) + 1
    return end + 0.5 * weight * (start - end) * cos_out

@HOOKS.register_module()
class CosineAnnealingLrUpdaterHook(LrUpdaterHook):

    def __init__(self, min_lr=None, min_lr_ratio=None, **kwargs):
        assert (min_lr is None) ^ (min_lr_ratio is None)
        self.min_lr = min_lr
        self.min_lr_ratio = min_lr_ratio
        super(CosineAnnealingLrUpdaterHook, self).__init__(**kwargs)

    def get_lr(self, runner, base_lr):
        if self.by_epoch:
            progress = runner.epoch
            max_progress = runner.max_epochs
        else:
            progress = runner.iter
            max_progress = runner.max_iters

        if self.min_lr_ratio is not None:
            target_lr = base_lr * self.min_lr_ratio
        else:
            target_lr = self.min_lr
        return annealing_cos(base_lr, target_lr, progress / max_progress)

从 get_lr 可以看到, 学习率变换周期有两种,epoch->max_epoch, 或者更大的 iter->max_iter, 后者表明一个 epoch 内不同 batch 的学习率可以不同, 因为没有什么理论, 所有这两种方式都行. 其中 base_lr 为初始学习率,target_lr 为学习率衰减的上界, 而当前学习率即为返回值.