mirrors/mmengine
1
0
Fork 0
mirror of https://github.com/open-mmlab/mmengine.git synced 2025-06-03 21:54:44 +08:00
Code Issues Projects Releases Wiki Activity

[Fix] Fix circle import in EMAHook (#523)

Browse Source
This commit is contained in:
Mashiro 2022-09-09 19:33:27 +08:00 committed by GitHub
parent 3ba50ce137
commit 0fb2b8ca8c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 6 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
mmengine/hooks/ema_hook.py
View File

@ -7,7 +7,6 @@ from typing import Dict, Optional
from mmengine.logging import print_log from mmengine.logging import print_log
from mmengine.model import is_model_wrapper from mmengine.model import is_model_wrapper
from mmengine.registry import HOOKS, MODELS from mmengine.registry import HOOKS, MODELS
from mmengine.runner.checkpoint import _load_checkpoint_to_model
from .hook import DATA_BATCH, Hook from .hook import DATA_BATCH, Hook
@ -172,6 +171,7 @@ class EMAHook(Hook):
Args: Args:
runner (Runner): The runner of the testing process. runner (Runner): The runner of the testing process.
""" """
from mmengine.runner.checkpoint import load_state_dict
if 'ema_state_dict' in checkpoint and runner._resume: if 'ema_state_dict' in checkpoint and runner._resume:
# The original model parameters are actually saved in ema # The original model parameters are actually saved in ema
# field swap the weights back to resume ema state. # field swap the weights back to resume ema state.
@ -186,7 +186,7 @@ class EMAHook(Hook):
'There is no `ema_state_dict` in checkpoint. ' 'There is no `ema_state_dict` in checkpoint. '
'`EMAHook` will make a copy of `state_dict` as the ' '`EMAHook` will make a copy of `state_dict` as the '
'initial `ema_state_dict`', 'current', logging.WARNING) 'initial `ema_state_dict`', 'current', logging.WARNING)
_load_checkpoint_to_model( load_state_dict(
self.ema_model.module, self.ema_model.module,
copy.deepcopy(checkpoint['state_dict']), copy.deepcopy(checkpoint['state_dict']),
strict=self.strict_load) strict=self.strict_load)

7
tests/test_hooks/test_ema_hook.py
View File

@ -60,7 +60,7 @@ class ToyModel3(BaseModel, ToyModel):
def __init__(self): def __init__(self):
super().__init__() super().__init__()
self.linear1 = nn.Linear(2, 2) self.linear = nn.Linear(2, 2)
def forward(self, *args, **kwargs): def forward(self, *args, **kwargs):
return super(BaseModel, self).forward(*args, **kwargs) return super(BaseModel, self).forward(*args, **kwargs)
@ -226,10 +226,11 @@ class TestEMAHook(TestCase):
test_evaluator=evaluator, test_evaluator=evaluator,
test_cfg=dict(), test_cfg=dict(),
work_dir=self.temp_dir.name, work_dir=self.temp_dir.name,
load_from=osp.join(self.temp_dir.name, 'epoch_2.pth'), load_from=osp.join(self.temp_dir.name,
'without_ema_state_dict.pth'),
default_hooks=dict(logger=None), default_hooks=dict(logger=None),
custom_hooks=[dict(type='EMAHook', strict_load=False)], custom_hooks=[dict(type='EMAHook', strict_load=False)],
experiment_name='test5') experiment_name='test5.1')
runner.test() runner.test()
# Test enable ema at 5 epochs. # Test enable ema at 5 epochs.