# Copyright (c) OpenMMLab. All rights reserved. import copy import os.path as osp import torch import torch.nn as nn from mmengine.config import ConfigDict from mmengine.hooks import EMAHook from mmengine.model import BaseModel, ExponentialMovingAverage from mmengine.registry import MODELS from mmengine.testing import RunnerTestCase, assert_allclose from mmengine.testing.runner_test_case import ToyModel class DummyWrapper(BaseModel): def __init__(self, model): super().__init__() if not isinstance(model, nn.Module): model = MODELS.build(model) self.module = model def forward(self, *args, **kwargs): return self.module(*args, **kwargs) class ToyModel2(ToyModel): def __init__(self): super().__init__() self.linear3 = nn.Linear(2, 1) def forward(self, *args, **kwargs): return super().forward(*args, **kwargs) class ToyModel3(ToyModel): def __init__(self): super().__init__() self.linear2 = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 1)) def forward(self, *args, **kwargs): return super().forward(*args, **kwargs) class TestEMAHook(RunnerTestCase): def setUp(self) -> None: MODELS.register_module(name='DummyWrapper', module=DummyWrapper) MODELS.register_module(name='ToyModel2', module=ToyModel2) MODELS.register_module(name='ToyModel3', module=ToyModel3) return super().setUp() def tearDown(self): MODELS.module_dict.pop('DummyWrapper') MODELS.module_dict.pop('ToyModel2') MODELS.module_dict.pop('ToyModel3') return super().tearDown() def test_init(self): EMAHook() with self.assertRaisesRegex(AssertionError, '`begin_iter` must'): EMAHook(begin_iter=-1) with self.assertRaisesRegex(AssertionError, '`begin_epoch` must'): EMAHook(begin_epoch=-1) with self.assertRaisesRegex(AssertionError, '`begin_iter` and `begin_epoch`'): EMAHook(begin_iter=1, begin_epoch=1) def _get_ema_hook(self, runner): for hook in runner.hooks: if isinstance(hook, EMAHook): return hook def test_before_run(self): cfg = copy.deepcopy(self.epoch_based_cfg) cfg.custom_hooks = [dict(type='EMAHook')] runner = self.build_runner(cfg) ema_hook = self._get_ema_hook(runner) ema_hook.before_run(runner) self.assertIsInstance(ema_hook.ema_model, ExponentialMovingAverage) self.assertIs(ema_hook.src_model, runner.model) def test_before_train(self): cfg = copy.deepcopy(self.epoch_based_cfg) cfg.custom_hooks = [ dict(type='EMAHook', begin_epoch=cfg.train_cfg.max_epochs - 1) ] runner = self.build_runner(cfg) ema_hook = self._get_ema_hook(runner) ema_hook.before_train(runner) cfg = copy.deepcopy(self.epoch_based_cfg) cfg.custom_hooks = [ dict(type='EMAHook', begin_epoch=cfg.train_cfg.max_epochs + 1) ] runner = self.build_runner(cfg) ema_hook = self._get_ema_hook(runner) with self.assertRaisesRegex(AssertionError, 'self.begin_epoch'): ema_hook.before_train(runner) cfg = copy.deepcopy(self.iter_based_cfg) cfg.custom_hooks = [ dict(type='EMAHook', begin_iter=cfg.train_cfg.max_iters + 1) ] runner = self.build_runner(cfg) ema_hook = self._get_ema_hook(runner) with self.assertRaisesRegex(AssertionError, 'self.begin_iter'): ema_hook.before_train(runner) def test_after_train_iter(self): cfg = copy.deepcopy(self.epoch_based_cfg) cfg.custom_hooks = [dict(type='EMAHook')] runner = self.build_runner(cfg) ema_hook = self._get_ema_hook(runner) ema_hook = self._get_ema_hook(runner) ema_hook.before_run(runner) ema_hook.before_train(runner) src_model = runner.model ema_model = ema_hook.ema_model with torch.no_grad(): for parameter in src_model.parameters(): parameter.data.copy_(torch.randn(parameter.shape)) ema_hook.after_train_iter(runner, 1) for src, ema in zip(src_model.parameters(), ema_model.parameters()): assert_allclose(src.data, ema.data) with torch.no_grad(): for parameter in src_model.parameters(): parameter.data.copy_(torch.randn(parameter.shape)) ema_hook.after_train_iter(runner, 1) for src, ema in zip(src_model.parameters(), ema_model.parameters()): self.assertFalse((src.data == ema.data).all()) def test_before_val_epoch(self): self._test_swap_parameters('before_val_epoch') def test_after_val_epoch(self): self._test_swap_parameters('after_val_epoch') def test_before_test_epoch(self): self._test_swap_parameters('before_test_epoch') def test_after_test_epoch(self): self._test_swap_parameters('after_test_epoch') def test_before_save_checkpoint(self): cfg = copy.deepcopy(self.epoch_based_cfg) runner = self.build_runner(cfg) checkpoint = dict(state_dict=ToyModel().state_dict()) ema_hook = EMAHook() ema_hook.before_run(runner) ema_hook.before_train(runner) ori_checkpoint = copy.deepcopy(checkpoint) ema_hook.before_save_checkpoint(runner, checkpoint) for key in ori_checkpoint['state_dict'].keys(): assert_allclose( ori_checkpoint['state_dict'][key].cpu(), checkpoint['ema_state_dict'][f'module.{key}'].cpu()) assert_allclose( ema_hook.ema_model.state_dict()[f'module.{key}'].cpu(), checkpoint['state_dict'][key].cpu()) def test_after_load_checkpoint(self): # Test load a checkpoint without ema_state_dict. cfg = copy.deepcopy(self.epoch_based_cfg) runner = self.build_runner(cfg) checkpoint = dict(state_dict=ToyModel().state_dict()) ema_hook = EMAHook() ema_hook.before_run(runner) ema_hook.before_train(runner) ema_hook.after_load_checkpoint(runner, checkpoint) for key in checkpoint['state_dict'].keys(): assert_allclose( checkpoint['state_dict'][key].cpu(), ema_hook.ema_model.state_dict()[f'module.{key}'].cpu()) # Test a warning should be raised when resuming from a checkpoint # without `ema_state_dict` runner._resume = True ema_hook.after_load_checkpoint(runner, checkpoint) with self.assertLogs(runner.logger, level='WARNING') as cm: ema_hook.after_load_checkpoint(runner, checkpoint) self.assertRegex(cm.records[0].msg, 'There is no `ema_state_dict`') # Check the weight of state_dict and ema_state_dict have been swapped. # when runner._resume is True runner._resume = True checkpoint = dict( state_dict=ToyModel().state_dict(), ema_state_dict=ExponentialMovingAverage(ToyModel()).state_dict()) ori_checkpoint = copy.deepcopy(checkpoint) ema_hook.after_load_checkpoint(runner, checkpoint) for key in ori_checkpoint['state_dict'].keys(): assert_allclose( ori_checkpoint['state_dict'][key].cpu(), ema_hook.ema_model.state_dict()[f'module.{key}'].cpu()) runner._resume = False ema_hook.after_load_checkpoint(runner, checkpoint) def test_with_runner(self): cfg = copy.deepcopy(self.epoch_based_cfg) cfg.custom_hooks = [ConfigDict(type='EMAHook')] runner = self.build_runner(cfg) ema_hook = self._get_ema_hook(runner) runner.train() self.assertTrue( isinstance(ema_hook.ema_model, ExponentialMovingAverage)) checkpoint = torch.load(osp.join(self.temp_dir.name, 'epoch_2.pth')) self.assertTrue('ema_state_dict' in checkpoint) self.assertTrue(checkpoint['ema_state_dict']['steps'] == 8) # load and testing cfg.load_from = osp.join(self.temp_dir.name, 'epoch_2.pth') runner = self.build_runner(cfg) runner.test() # with model wrapper cfg.model = ConfigDict(type='DummyWrapper', model=cfg.model) runner = self.build_runner(cfg) runner.test() # Test load checkpoint without ema_state_dict checkpoint = torch.load(osp.join(self.temp_dir.name, 'epoch_2.pth')) checkpoint.pop('ema_state_dict') torch.save(checkpoint, osp.join(self.temp_dir.name, 'without_ema_state_dict.pth')) cfg.load_from = osp.join(self.temp_dir.name, 'without_ema_state_dict.pth') runner = self.build_runner(cfg) runner.test() # Test does not load checkpoint strictly (different name). # Test load checkpoint without ema_state_dict cfg.model = ConfigDict(type='ToyModel2') cfg.custom_hooks = [ConfigDict(type='EMAHook', strict_load=False)] runner = self.build_runner(cfg) runner.test() # Test does not load ckpt strictly (different weight size). # Test load checkpoint without ema_state_dict cfg.model = ConfigDict(type='ToyModel3') runner = self.build_runner(cfg) runner.test() # Test enable ema at 5 epochs. cfg.train_cfg.max_epochs = 10 cfg.custom_hooks = [ConfigDict(type='EMAHook', begin_epoch=5)] runner = self.build_runner(cfg) runner.train() state_dict = torch.load( osp.join(self.temp_dir.name, 'epoch_4.pth'), map_location='cpu') self.assertIn('ema_state_dict', state_dict) for k, v in state_dict['state_dict'].items(): assert_allclose(v, state_dict['ema_state_dict']['module.' + k]) # Test enable ema at 5 iterations. cfg = copy.deepcopy(self.iter_based_cfg) cfg.train_cfg.val_interval = 1 cfg.custom_hooks = [ConfigDict(type='EMAHook', begin_iter=5)] cfg.default_hooks.checkpoint.interval = 1 runner = self.build_runner(cfg) runner.train() state_dict = torch.load( osp.join(self.temp_dir.name, 'iter_4.pth'), map_location='cpu') self.assertIn('ema_state_dict', state_dict) for k, v in state_dict['state_dict'].items(): assert_allclose(v, state_dict['ema_state_dict']['module.' + k]) state_dict = torch.load( osp.join(self.temp_dir.name, 'iter_5.pth'), map_location='cpu') self.assertIn('ema_state_dict', state_dict) def _test_swap_parameters(self, func_name, *args, **kwargs): cfg = copy.deepcopy(self.epoch_based_cfg) cfg.custom_hooks = [dict(type='EMAHook')] runner = self.build_runner(cfg) ema_hook = self._get_ema_hook(runner) runner.train() with torch.no_grad(): for parameter in ema_hook.src_model.parameters(): parameter.data.copy_(torch.randn(parameter.shape)) src_model = copy.deepcopy(runner.model) ema_model = copy.deepcopy(ema_hook.ema_model) func = getattr(ema_hook, func_name) func(runner, *args, **kwargs) swapped_src = ema_hook.src_model swapped_ema = ema_hook.ema_model for src, ema, swapped_src, swapped_ema in zip( src_model.parameters(), ema_model.parameters(), swapped_src.parameters(), swapped_ema.parameters()): self.assertTrue((src.data == swapped_ema.data).all()) self.assertTrue((ema.data == swapped_src.data).all())