MMSelfSup在训练过程中调用当前Epoch

本文参考了这篇博客：mmdet在训练过程中调用epoch_mmdet custom_hooks-CSDN博客

由于MMSelfSup训练过程都封装到runner中，在训练过程中是无法直接调用epoch。而我的模型在损失函数计算时，正好有个参数是与模型训练的epoch有关的（）。因此，在每次训练时需要调用模型当前训练的epoch。

首先，自定义一个hook函数，命名为SetEpochInfoHook，位置在mmselfsup-main/mmselfsup/engine/hooks

from mmengine.hooks import Hook

from mmselfsup.registry import HOOKS
from mmengine.model import is_model_wrapper

@HOOKS.register_module()
class SetEpochInfoHook(Hook):
    """Set runner's epoch information to the model."""

    def before_train_epoch(self, runner):
        epoch = runner.epoch+1
        model = runner.model
        if is_model_wrapper(model):
            model = model.module
        model.set_epoch(epoch)

在__init__函数中把SetEpochInfoHook添加进去。

然后，在cfg配置文件中，添加以下代码，自定义一个hook，这个hook类型是我们自定义的hook。

custom_hooks = [dict(type='SetEpochInfoHook')] # my own

在自定义的hook中，使用了一个model的set_epoch的方法。因此，在model的代码中，把这个方法需要具体实现下。在运行SetEpochInfoHook时，会在程序中自动运行model的set_epoch方法。

我使用的模型的算法代码在mmselfsup-main/mmselfsup/models/algorithms/simclr_hash.py中，因此，我又给模型定义了一个函数。

   def set_epoch(self,epoch):
        self.epoch=epoch

在调用时，直接用self.epoch就可以获取模型当前的epoch，实现了调用。