[Fix] Fix ema hook and add unit test (#327)

* Fix ema hook and add unit test

* save state_dict of ema.module

save state_dict of ema.module

* replace warning.warn with MMLogger.warn

* fix as comment

* fix bug

* fix bug

mmengine/hooks/ema_hook.py

 # Copyright (c) OpenMMLab. All rights reserved.
+import copy
 import itertools
+import logging
 from typing import Dict, Optional
 
+from mmengine.logging import print_log
 from mmengine.model import is_model_wrapper
 from mmengine.registry import HOOKS, MODELS
 from .hook import DATA_BATCH, Hook
@@ -80,10 +83,21 @@ class EMAHook(Hook):
 
     def after_load_checkpoint(self, runner, checkpoint: dict) -> None:
         """Resume ema parameters from checkpoint."""
-        # The original model parameters are actually saved in ema field.
-        # swap the weights back to resume ema state.
-        self._swap_ema_state_dict(checkpoint)
-        self.ema_model.load_state_dict(checkpoint['ema_state_dict'])
+
+        if 'ema_state_dict' in checkpoint:
+            # The original model parameters are actually saved in ema field.
+            # swap the weights back to resume ema state.
+            self._swap_ema_state_dict(checkpoint)
+            self.ema_model.load_state_dict(checkpoint['ema_state_dict'])
+
+        # Support load checkpoint without ema state dict.
+        else:
+            print_log(
+                'There is no `ema_state_dict` in checkpoint. '
+                '`EMAHook` will make a copy of `state_dict` as the '
+                'initial `ema_state_dict`', 'current', logging.WARNING)
+            self.ema_model.module.load_state_dict(
+                copy.deepcopy(checkpoint['state_dict']))
 
     def _swap_ema_parameters(self) -> None:
         """Swap the parameter of model with ema_model."""

tests/test_hook/test_ema_hook.py

@@ -149,3 +149,25 @@ class TestEMAHook(TestCase):
             custom_hooks=[dict(type='EMAHook')],
             experiment_name='test3')
         runner.test()
+
+        # Test load checkpoint without ema_state_dict
+        ckpt = torch.load(osp.join(self.temp_dir.name, 'epoch_2.pth'))
+        ckpt.pop('ema_state_dict')
+        torch.save(ckpt,
+                   osp.join(self.temp_dir.name, 'without_ema_state_dict.pth'))
+        runner = Runner(
+            model=DummyWrapper(ToyModel()),
+            test_dataloader=dict(
+                dataset=dict(type='DummyDataset'),
+                sampler=dict(type='DefaultSampler', shuffle=True),
+                batch_size=3,
+                num_workers=0),
+            test_evaluator=evaluator,
+            test_cfg=dict(),
+            work_dir=self.temp_dir.name,
+            load_from=osp.join(self.temp_dir.name,
+                               'without_ema_state_dict.pth'),
+            default_hooks=dict(logger=None),
+            custom_hooks=[dict(type='EMAHook')],
+            experiment_name='test4')
+        runner.test()