mirror of
https://github.com/open-mmlab/mmclassification.git
synced 2025-06-03 21:53:55 +08:00
08dc8c75d3
* remove basehead * add moco series * add byol simclr simsiam * add ut * update configs * add simsiam hook * add and refactor beit * update ut * add cae * update extract_feat * refactor cae * add mae * refactor data preprocessor * update heads * add maskfeat * add milan * add simmim * add mixmim * fix lint * fix ut * fix lint * add eva * add densecl * add barlowtwins * add swav * fix lint * update readtherdocs rst * update docs * update * Decrease UT memory usage * Fix docstring * update DALLEEncoder * Update model docs * refactor dalle encoder * update docstring * fix ut * fix config error * add val_cfg and test_cfg * refactor clip generator * fix lint * pass check * fix ut * add lars * update type of BEiT in configs * Use MMEngine style momentum in EMA. * apply mmpretrain solarize --------- Co-authored-by: mzr1996 <mzr1996@163.com>
43 lines
1.5 KiB
Python
43 lines
1.5 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
from typing import Optional, Sequence
|
|
|
|
from mmengine.hooks import Hook
|
|
|
|
from mmpretrain.registry import HOOKS
|
|
from mmpretrain.utils import get_ori_model
|
|
|
|
|
|
@HOOKS.register_module()
|
|
class DenseCLHook(Hook):
|
|
"""Hook for DenseCL.
|
|
|
|
This hook includes ``loss_lambda`` warmup in DenseCL.
|
|
Borrowed from the authors' code: `<https://github.com/WXinlong/DenseCL>`_.
|
|
|
|
Args:
|
|
start_iters (int): The number of warmup iterations to set
|
|
``loss_lambda=0``. Defaults to 1000.
|
|
"""
|
|
|
|
def __init__(self, start_iters: int = 1000) -> None:
|
|
self.start_iters = start_iters
|
|
|
|
def before_train(self, runner) -> None:
|
|
"""Obtain ``loss_lambda`` from algorithm."""
|
|
assert hasattr(get_ori_model(runner.model), 'loss_lambda'), \
|
|
"The runner must have attribute \"loss_lambda\" in DenseCL."
|
|
self.loss_lambda = get_ori_model(runner.model).loss_lambda
|
|
|
|
def before_train_iter(self,
|
|
runner,
|
|
batch_idx: int,
|
|
data_batch: Optional[Sequence[dict]] = None) -> None:
|
|
"""Adjust ``loss_lambda`` every train iter."""
|
|
assert hasattr(get_ori_model(runner.model), 'loss_lambda'), \
|
|
"The runner must have attribute \"loss_lambda\" in DenseCL."
|
|
cur_iter = runner.iter
|
|
if cur_iter >= self.start_iters:
|
|
get_ori_model(runner.model).loss_lambda = self.loss_lambda
|
|
else:
|
|
get_ori_model(runner.model).loss_lambda = 0.
|