RSPrompter / mmpretrain /engine /hooks /retriever_hooks.py
KyanChen's picture
Upload 303 files
4d0eb62
raw
history blame
1.04 kB
# Copyright (c) OpenMMLab. All rights reserved
import warnings
from mmengine.hooks import Hook
from mmengine.model import is_model_wrapper
from mmpretrain.models import BaseRetriever
from mmpretrain.registry import HOOKS
@HOOKS.register_module()
class PrepareProtoBeforeValLoopHook(Hook):
"""The hook to prepare the prototype in retrievers.
Since the encoders of the retriever changes during training, the prototype
changes accordingly. So the `prototype_vecs` needs to be regenerated before
validation loop.
"""
def before_val(self, runner) -> None:
model = runner.model
if is_model_wrapper(model):
model = model.module
if isinstance(model, BaseRetriever):
if hasattr(model, 'prepare_prototype'):
model.prepare_prototype()
else:
warnings.warn(
'Only the `mmpretrain.models.retrievers.BaseRetriever` '
'can execute `PrepareRetrieverPrototypeHook`, but got '
f'`{type(model)}`')