mmdet在训练过程中调用epoch

由于mmdet训练过程都封装到runner中,在训练过程中是无法直接调用epoch。本文主要参考了

https://github.com/open-mmlab/mmdetection/issues/7425

具体实现

  1. 定义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文件中。

  1. 修改配置文件
    在配置文件中,添加
# custom hooks
custom_hooks = [dict(type='SetEpochInfoHook')]
  1. 注册hook
    mmdet/apis/train.py中,添加如下代码
    mmdet在训练过程中调用epoch_第1张图片
    SetEpochInfoHook就是刚才新增的类。

  2. 在detector文件中修改
    我是在two_stage.py中,添加set_epoch函数。
    mmdet在训练过程中调用epoch_第2张图片
    5.调用方法
    这样就可以在训练过程中,获取当前的epoch。
    mmdet在训练过程中调用epoch_第3张图片
    print的结果是1。

你可能感兴趣的:(mmdetion,python,人工智能)