# 迁移 MMCV 钩子到 MMEngine
## 简介
由于架构设计的更新和用户需求的不断增加,MMCV 的钩子(Hook)点位已经满足不了需求,因此在 MMEngine 中对钩子点位进行了重新设计以及对钩子的功能做了调整。在开始迁移前,阅读[钩子的设计](../design/hook.md)会很有帮助。
本文对比 [MMCV v1.6.0](https://github.com/open-mmlab/mmcv/tree/v1.6.0) 和 [MMEngine v0.5.0](https://github.com/open-mmlab/mmengine/tree/v0.5.0) 的钩子在功能、点位、用法和实现上的差异。
## 功能差异
|
MMCV |
MMEngine |
反向传播以及梯度更新 |
OptimizerHook |
将反向传播以及梯度更新的操作抽象成 OptimWrapper 而不是钩子 |
GradientCumulativeOptimizerHook |
学习率调整 |
LrUpdaterHook |
ParamSchdulerHook 以及 _ParamScheduler 的子类完成优化器超参的调整 |
动量调整 |
MomentumUpdaterHook |
按指定间隔保存权重 |
CheckpointHook |
CheckpointHook 除了保存权重,还有保存最优权重的功能,而 EvalHook 的模型评估功能则交由 ValLoop 或 TestLoop 完成 |
模型评估并保存最优模型 |
EvalHook |
打印日志 |
LoggerHook 及其子类实现打印日志、保存日志以及可视化功能 |
LoggerHook |
可视化 |
NaiveVisualizationHook |
添加运行时信息 |
RuntimeInfoHook |
模型参数指数滑动平均 |
EMAHook |
EMAHook |
确保分布式 Sampler 的 shuffle 生效 |
DistSamplerSeedHook |
DistSamplerSeedHook |
同步模型的 buffer |
SyncBufferHook |
SyncBufferHook |
PyTorch CUDA 缓存清理 |
EmptyCacheHook |
EmptyCacheHook |
统计迭代耗时 |
IterTimerHook |
IterTimerHook |
分析训练时间的瓶颈 |
ProfilerHook |
暂未提供 |
提供注册方法给钩子点位的功能 |
ClosureHook |
暂未提供 |
## 点位差异
|
MMCV |
MMEngine |
全局位点 |
执行前 |
before_run |
before_run |
执行后 |
after_run |
after_run |
Checkpoint 相关 |
加载 checkpoint 后 |
无 |
after_load_checkpoint |
保存 checkpoint 前 |
无 |
before_save_checkpoint |
训练相关 |
训练前触发 |
无 |
before_train |
训练后触发 |
无 |
after_train |
每个 epoch 前 |
before_train_epoch |
before_train_epoch |
每个 epoch 后 |
after_train_epoch |
after_train_epoch |
每次迭代前 |
before_train_iter |
before_train_iter,新增 batch_idx 和 data_batch 参数 |
每次迭代后 |
after_train_iter |
after_train_iter,新增 batch_idx、data_batch 和 outputs 参数 |
验证相关 |
验证前触发 |
无 |
before_val |
验证后触发 |
无 |
after_val |
每个 epoch 前 |
before_val_epoch |
before_val_epoch |
每个 epoch 后 |
after_val_epoch |
after_val_epoch |
每次迭代前 |
before_val_iter |
before_val_iter,新增 batch_idx 和 data_batch 参数 |
每次迭代后 |
after_val_iter |
after_val_iter,新增 batch_idx、data_batch 和 outputs 参数 |
测试相关 |
测试前触发 |
无 |
before_test |
测试后触发 |
无 |
after_test |
每个 epoch 前 |
无 |
before_test_epoch |
每个 epoch 后 |
无 |
after_test_epoch |
每次迭代前 |
无 |
before_test_iter,新增 batch_idx 和 data_batch 参数 |
每次迭代后 |
无 |
after_test_iter,新增 batch_idx、data_batch 和 outputs 参数 |
## 用法差异
在 MMCV 中,将钩子注册到执行器(Runner),需调用执行器的 `register_training_hooks` 方法往执行器注册钩子,而在 MMEngine 中,可以通过参数传递给执行器的初始化方法进行注册。
- MMCV
```python
model = ResNet18()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
lr_config = dict(policy='step', step=[2, 3])
optimizer_config = dict(grad_clip=None)
checkpoint_config = dict(interval=5)
log_config = dict(interval=100, hooks=[dict(type='TextLoggerHook')])
custom_hooks = [dict(type='NumClassCheckHook')]
runner = EpochBasedRunner(
model=model,
optimizer=optimizer,
work_dir='./work_dir',
max_epochs=3,
xxx,
)
runner.register_training_hooks(
lr_config=lr_config,
optimizer_config=optimizer_config,
checkpoint_config=checkpoint_config,
log_config=log_config,
custom_hooks_config=custom_hooks,
)
runner.run([trainloader], [('train', 1)])
```
- MMEngine
```python
model=ResNet18()
optim_wrapper=dict(
type='OptimizerWrapper',
optimizer=dict(type='SGD', lr=0.001, momentum=0.9))
param_scheduler = dict(type='MultiStepLR', milestones=[2, 3]),
default_hooks = dict(
logger=dict(type='LoggerHook'),
param_scheduler=dict(type='ParamSchedulerHook'),
checkpoint=dict(type='CheckpointHook', interval=5),
)
custom_hooks = [dict(type='NumClassCheckHook')]
runner = Runner(
model=model,
work_dir='./work_dir',
optim_wrapper=optim_wrapper,
param_scheduler=param_scheduler,
train_cfg=dict(by_epoch=True, max_epochs=3),
default_hooks=default_hooks,
custom_hooks=custom_hooks,
xxx,
)
runner.train()
```
MMEngine 钩子的更多用法请参考[钩子的用法](../tutorials/hook.md)。
## 实现差异
以 `CheckpointHook` 为例,MMEngine 的 [CheckpointHook](https://github.com/open-mmlab/mmengine/blob/main/mmengine/hooks/checkpoint_hook.py) 相比 MMCV 的 [CheckpointHook](https://github.com/open-mmlab/mmcv/blob/v1.6.0/mmcv/runner/hooks/checkpoint.py)(新增保存最优权重的功能,在 MMCV 中,保存最优权重的功能由 EvalHook 提供),因此,它需要实现 `after_val_epoch` 点位。
- MMCV
```python
class CheckpointHook(Hook):
def before_run(self, runner):
"""初始化 out_dir 和 file_client 属性"""
def after_train_epoch(self, runner):
"""同步 buffer 和保存权重,用于以 epoch 为单位训练的任务"""
def after_train_iter(self, runner):
"""同步 buffer 和保存权重,用于以 iteration 为单位训练的任务"""
```
- MMEngine
```python
class CheckpointHook(Hook):
def before_run(self, runner):
"""初始化 out_dir 和 file_client 属性"""
def after_train_epoch(self, runner):
"""同步 buffer 和保存权重,用于以 epoch 为单位训练的任务"""
def after_train_iter(self, runner, batch_idx, data_batch, outputs):
"""同步 buffer 和保存权重,用于以 iteration 为单位训练的任务"""
def after_val_epoch(self, runner, metrics):
"""根据 metrics 保存最优权重"""
```