mmdetection-3.x学习笔记——MMEngin核心组件分析（一）：Hook
Hook编程是一种编程模式，是指在程序的一个或者多个位置设置位点（挂载点），当程序运行至某个位点时，会自动调用运行时注册到位点的所有方法。在 python 中由于函数是一等公民，实现 hook 机制其实只需要传入一个函数即可，在该函数中我们可以获取到内部信息，可以修改或者访问该内容，从而实现一些类似黑科技一般的功能，如下是一个简单的Hook，打印内部变量，
def hook(d):
print(d)
def add(a,b,c,hook_fn=None)
sum1=a+b
if hook_fn is not None:
hook_fn(sum1)
return sum1+c
# 调用
add(1,2,3,hook)
Hook配合Runner可以管理整个模型训练流程。
在Runner（/mmengine/runner/runner.py）类中，维护了一个Hook池，Runner初始化时会把用到的Hook（default_hooks, custom_hooks）都注册到Hook池中，
@RUNNERS.register_module()
class Runner:
def __init__(self,
...,
default_hooks: Optional[Dict[str, Union[Hook, Dict]]] = None,
custom_hooks: Optional[List[Union[Hook, Dict]]] = None,
...
):
.......
self._hooks: List[Hook] = [] # Hook池初始化
# register hooks to `self._hooks`
self.register_hooks(default_hooks, custom_hooks) # 调用register_hooks()函数把default_hooks, custom_hooks都注册到Hook池中
# log hooks information
self.logger.info(f'Hooks will be executed in the following '
f'order:\n{self.get_hooks_info()}')
因此，有default_hooks和custom_hooks两大类，所以用register_hooks()函数再封装一层，register_hooks()函数内部调用register_default_hooks()和register_custom_hooks()函数，
def register_hooks(
self,
default_hooks: Optional[Dict[str, Union[Hook, Dict]]] = None,
custom_hooks: Optional[List[Union[Hook, Dict]]] = None) -> None:
"""Register default hooks and custom hooks into hook list.
Args:
default_hooks (dict[str, dict] or dict[str, Hook], optional): Hooks
to execute default actions like updating model parameters and
saving checkpoints. Defaults to None.
custom_hooks (list[dict] or list[Hook], optional): Hooks to execute
custom actions like visualizing images processed by pipeline.
Defaults to None.
"""
self.register_default_hooks(default_hooks)
if custom_hooks is not None:
self.register_custom_hooks(custom_hooks)
register_default_hooks(）函数定义默认的Hooks，并调用register_hook()函数注册到Hook池中，
def register_default_hooks(
self,
hooks: Optional[Dict[str, Union[Hook, Dict]]] = None) -> None:
"""Register default hooks into hook list.
``hooks`` will be registered into runner to execute some default
actions like updating model parameters or saving checkpoints.
Default hooks and their priorities:
+----------------------+-------------------------+
| Hooks | Priority |
+======================+=========================+
| RuntimeInfoHook | VERY_HIGH (10) |
+----------------------+-------------------------+
| IterTimerHook | NORMAL (50) |
+----------------------+-------------------------+
| DistSamplerSeedHook | NORMAL (50) |
+----------------------+-------------------------+
| LoggerHook | BELOW_NORMAL (60) |
+----------------------+-------------------------+
| ParamSchedulerHook | LOW (70) |
+----------------------+-------------------------+
| CheckpointHook | VERY_LOW (90) |
+----------------------+-------------------------+
If ``hooks`` is None, above hooks will be registered by
default::
default_hooks = dict(
runtime_info=dict(type='RuntimeInfoHook'),
timer=dict(type='IterTimerHook'),
sampler_seed=dict(type='DistSamplerSeedHook'),
logger=dict(type='LoggerHook'),
param_scheduler=dict(type='ParamSchedulerHook'),
checkpoint=dict(type='CheckpointHook', interval=1),
)
If not None, ``hooks`` will be merged into ``default_hooks``.
If there are None value in default_hooks, the corresponding item will
be popped from ``default_hooks``::
hooks = dict(timer=None)
The final registered default hooks will be :obj:`RuntimeInfoHook`,
:obj:`DistSamplerSeedHook`, :obj:`LoggerHook`,
:obj:`ParamSchedulerHook` and :obj:`CheckpointHook`.
Args:
hooks (dict[str, Hook or dict], optional): Default hooks or configs
to be registered.
"""
default_hooks: dict = dict( # 定义默认的Hook
runtime_info=dict(type='RuntimeInfoHook'),
timer=dict(type='IterTimerHook'),
sampler_seed=dict(type='DistSamplerSeedHook'),
logger=dict(type='LoggerHook'),
param_scheduler=dict(type='ParamSchedulerHook'),
checkpoint=dict(type='CheckpointHook', interval=1),
)
if hooks is not None:
for name, hook in hooks.items():
if name in default_hooks and hook is None: # 移除不想要的Hook，假如不想要timer这个hook，则把hooks = dict(timer=None)作为参数传入
# remove hook from _default_hooks
default_hooks.pop(name)
else:
assert hook is not None
default_hooks[name] = hook # hook不为空，且hook不在default_hooks内，即将该hook添加到defau_hooks中
for hook in default_hooks.values():
self.register_hook(hook) # 真正的注册函数
Hook注册函数，将hook按照优先级插入到Hook池中，
def register_hook(
self,
hook: Union[Hook, Dict],
priority: Optional[Union[str, int, Priority]] = None) -> None:
"""Register a hook into the hook list.
The hook will be inserted into a priority queue, with the specified
priority (See :class:`Priority` for details of priorities).
For hooks with the same priority, they will be triggered in the same
order as they are registered.
Priority of hook will be decided with the following priority:
- ``priority`` argument. If ``priority`` is given, it will be priority
of hook.
- If ``hook`` argument is a dict and ``priority`` in it, the priority
will be the value of ``hook['priority']``.
- If ``hook`` argument is a dict but ``priority`` not in it or ``hook``
is an instance of ``hook``, the priority will be ``hook.priority``.
Args:
hook (:obj:`Hook` or dict): The hook to be registered.
priority (int or str or :obj:`Priority`, optional): Hook priority.
Lower value means higher priority.
"""
if not isinstance(hook, (Hook, dict)):
raise TypeError(
f'hook should be an instance of Hook or dict, but got {hook}')
_priority = None
if isinstance(hook, dict): # 传入的是hook配置词典
if 'priority' in hook:
_priority = hook.pop('priority') # 获取hook配置词典中指定的优先级
hook_obj = HOOKS.build(hook) # 构建hook
else:
hook_obj = hook # 传入的是hook对象
if priority is not None: # 调用注册函数时指定优先级
hook_obj.priority = priority
elif _priority is not None:
hook_obj.priority = _priority
inserted = False
for i in range(len(self._hooks) - 1, -1, -1):
if get_priority(hook_obj.priority) >= get_priority(
self._hooks[i].priority): # 比较要注册的hook的优先级与hook池中的优先级，找到正确的位置并插入
self._hooks.insert(i + 1, hook_obj)
inserted = True
break
if not inserted:
self._hooks.insert(0, hook_obj) # 优先级最高
Runner在模型的训练、测试流程中插入位点，并在对应位点实现调用 hook 的抽象逻辑。如在Runner的train函数中插入了'before_run'和'after_run'两个位点，
def train(self) -> nn.Module:
......
self.call_hook('before_run')
......
model = self.train_loop.run() # type: ignore
self.call_hook('after_run')
return model
在训练循环中插入了'before_train'和'after_train'两个位点，
def run(self) -> torch.nn.Module:
"""Launch training."""
self.runner.call_hook('before_train')
while self._epoch < self._max_epochs:
self.run_epoch()
self._decide_current_val_interval()
if (self.runner.val_loop is not None
and self._epoch >= self.val_begin
and self._epoch % self.val_interval == 0):
self.runner.val_loop.run()
self.runner.call_hook('after_train')
return self.runner.model
在run_epoch()中插入了'before_train_epoch'和'after_train_epoch'两个位点
def run_epoch(self) -> None:
"""Iterate one epoch."""
self.runner.call_hook('before_train_epoch')
self.runner.model.train()
for idx, data_batch in enumerate(self.dataloader):
self.run_iter(idx, data_batch)
self.runner.call_hook('after_train_epoch')
self._epoch += 1
在run_iter()中插入了'before_train_iter'和'after_train_iter'两个位点
def run_iter(self, idx, data_batch: Sequence[dict]) -> None:
"""Iterate one min-batch.
Args:
data_batch (Sequence[dict]): Batch of data from dataloader.
"""
self.runner.call_hook(
'before_train_iter', batch_idx=idx, data_batch=data_batch)
# Enable gradient accumulation mode and avoid unnecessary gradient
# synchronization during gradient accumulation process.
# outputs should be a dict of loss.
outputs = self.runner.model.train_step(
data_batch, optim_wrapper=self.runner.optim_wrapper)
self.runner.call_hook(
'after_train_iter',
batch_idx=idx,
data_batch=data_batch,
outputs=outputs)
self._iter += 1
Hook中一共有 22 个位点，在伪代码位置如下
def main():
...
call_hooks('before_run', hooks) # 任务开始前执行的逻辑
call_hooks('after_load_checkpoint', hooks) # 加载权重后执行的逻辑
call_hooks('before_train', hooks) # 训练开始前执行的逻辑
for i in range(max_epochs):
call_hooks('before_train_epoch', hooks) # 遍历训练数据集前执行的逻辑
for inputs, labels in train_dataloader:
call_hooks('before_train_iter', hooks) # 模型前向计算前执行的逻辑
outputs = net(inputs)
loss = criterion(outputs, labels)
call_hooks('after_train_iter', hooks) # 模型前向计算后执行的逻辑
loss.backward()
optimizer.step()
call_hooks('after_train_epoch', hooks) # 遍历完训练数据集后执行的逻辑
call_hooks('before_val_epoch', hooks) # 遍历验证数据集前执行的逻辑
with torch.no_grad():
for inputs, labels in val_dataloader:
call_hooks('before_val_iter', hooks) # 模型前向计算前执行
outputs = net(inputs)
loss = criterion(outputs, labels)
call_hooks('after_val_iter', hooks) # 模型前向计算后执行
call_hooks('after_val_epoch', hooks) # 遍历完验证数据集前执行
call_hooks('before_save_checkpoint', hooks) # 保存权重前执行的逻辑
call_hooks('after_train', hooks) # 训练结束后执行的逻辑
call_hooks('before_test_epoch', hooks) # 遍历测试数据集前执行的逻辑
with torch.no_grad():
for inputs, labels in test_dataloader:
call_hooks('before_test_iter', hooks) # 模型前向计算后执行的逻辑
outputs = net(inputs)
accuracy = ...
call_hooks('after_test_iter', hooks) # 遍历完成测试数据集后执行的逻辑
call_hooks('after_test_epoch', hooks) # 遍历完测试数据集后执行
call_hooks('after_run', hooks) # 任务结束后执行的逻辑
在每个位点通过call_hook()调用所有hooks在该位点对应的函数，它会遍历整个hook池，然后执行每个hooks对应该位点的函数，如在'before_run'这个位点，则执行每个hook的before_run()函数。call_hook()的实现逻辑如下，
def call_hook(self, fn_name: str, **kwargs) -> None:
"""Call all hooks.
Args:
fn_name (str): The function name in each hook to be called, such as
"before_train_epoch".
**kwargs: Keyword arguments passed to hook.
"""
for hook in self._hooks: # 遍历hook池
# support adding additional custom hook methods
if hasattr(hook, fn_name): # 如果该hook有该位点对应的函数
try:
getattr(hook, fn_name)(self, **kwargs) # 则执行
except TypeError as e:
raise TypeError(f'{e} in {hook}') from None
下面是YOLOXModeSwitchHook，它的作用是训练的最后15个epoch关闭mosaic 和 mixup数据增广，并启用L1损失函数对回归分支微调，所以在每个epoch开始前需要判断当前是否是最后15个epoch，如果是则执行对应操作。因此，该hook只需要before_train_epoch()这一个函数。
# Copyright (c) OpenMMLab. All rights reserved.
from typing import Sequence
from mmengine.hooks import Hook
from mmengine.model import is_model_wrapper
from mmdet.registry import HOOKS
@HOOKS.register_module()
class YOLOXModeSwitchHook(Hook):
"""Switch the mode of YOLOX during training.
This hook turns off the mosaic and mixup data augmentation and switches
to use L1 loss in bbox_head.
Args:
num_last_epochs (int): The number of latter epochs in the end of the
training to close the data augmentation and switch to L1 loss.
Defaults to 15.
skip_type_keys (Sequence[str], optional): Sequence of type string to be
skip pipeline. Defaults to ('Mosaic', 'RandomAffine', 'MixUp').
"""
def __init__(
self,
num_last_epochs: int = 15,
skip_type_keys: Sequence[str] = ('Mosaic', 'RandomAffine', 'MixUp')
) -> None:
self.num_last_epochs = num_last_epochs
self.skip_type_keys = skip_type_keys
self._restart_dataloader = False
def before_train_epoch(self, runner) -> None:
"""Close mosaic and mixup augmentation and switches to use L1 loss."""
epoch = runner.epoch
train_loader = runner.train_dataloader
model = runner.model
# TODO: refactor after mmengine using model wrapper
if is_model_wrapper(model):
model = model.module
if (epoch + 1) == runner.max_epochs - self.num_last_epochs:
runner.logger.info('No mosaic and mixup aug now!')
# The dataset pipeline cannot be updated when persistent_workers
# is True, so we need to force the dataloader's multi-process
# restart. This is a very hacky approach.
train_loader.dataset.update_skip_type_keys(self.skip_type_keys)
if hasattr(train_loader, 'persistent_workers'
) and train_loader.persistent_workers is True:
train_loader._DataLoader__initialized = False
train_loader._iterator = None
self._restart_dataloader = True
runner.logger.info('Add additional L1 loss now!')
model.bbox_head.use_l1 = True
else:
# Once the restart is complete, we need to restore
# the initialization flag.
if self._restart_dataloader:
train_loader._DataLoader__initialized = True