由于mmdet训练过程都封装到runner中,在训练过程中是无法直接调用epoch。本文主要参考了
https://github.com/open-mmlab/mmdetection/issues/7425
具体实现
- 定义hook类
在mmdet/core/hooks
中新增文件set_epoch_info_hook.py
,内容如下:
# Copyright (c) OpenMMLab. All rights reserved.
from mmcv.parallel import is_module_wrapper
from mmcv.runner import HOOKS, Hook
@HOOKS.register_module()
class SetEpochInfoHook(Hook):
"""Set runner's epoch information to the model."""
def before_train_epoch(self, runner):
epoch = runner.epoch
model = runner.model
if is_module_wrapper(model):
model = model.module
model.set_epoch(epoch)
并把该类添加到对应的__init__.py文件中。
- 修改配置文件
在配置文件中,添加
# custom hooks
custom_hooks = [dict(type='SetEpochInfoHook')]
-
注册hook
在mmdet/apis/train.py
中,添加如下代码
SetEpochInfoHook就是刚才新增的类。 -
在detector文件中修改
我是在two_stage.py
中,添加set_epoch函数。
5.调用方法
这样就可以在训练过程中,获取当前的epoch。
print的结果是1。