Add load rewriter (#43)

* Add load rewriter

* Fix resume

* Fix typo

* Fix test code

Author: yhna940

configs/mmtune/_base_/context/train.py

@@ -12,6 +12,7 @@
         src_key='searched_cfg',
         dst_key='base_cfg',
         key='cfg'),
+    dict(type='ResumeFromCkpt'),
     dict(
         type='CustomHookRegister',
         key='cfg',

mmtune/mm/context/rewriters/__init__.py

@@ -6,9 +6,10 @@
 from .patch import BatchConfigPatcher, SequeunceConfigPatcher
 from .path import AppendTrialIDtoPath
 from .register import CustomHookRegister
+from .resume import ResumeFromCkpt
 
 __all__ = [
     'BaseRewriter', 'REWRITERS', 'build_rewriter', 'Dump', 'MergeConfig',
     'AppendTrialIDtoPath', 'BatchConfigPatcher', 'SequeunceConfigPatcher',
-    'CustomHookRegister', 'InstantiateCfg'
+    'CustomHookRegister', 'InstantiateCfg', 'ResumeFromCkpt'
 ]

mmtune/mm/context/rewriters/resume.py

@@ -0,0 +1,30 @@
+from typing import Dict
+
+from .base import BaseRewriter
+from .builder import REWRITERS
+
+
+@REWRITERS.register_module()
+class ResumeFromCkpt(BaseRewriter):
+    """Specifies the checkpoint for resuming training."""
+
+    def __init__(self, arg_name: str = 'resume_from') -> None:
+        """Initialize the rewriter.
+
+        Args:
+            key (str): The key where the instantiated cfg is stored.
+            arg_name (str): The key in the argparse namespace.
+        """
+        self.arg_name = arg_name
+
+    def __call__(self, context: Dict) -> Dict:
+        """Set with checkpoints specified by Ray.
+
+        Args:
+            context (Dict): The context to be rewritten.
+        Returns:
+            Dict: The context after rewriting.
+        """
+        setattr(
+            context.get('args'), self.arg_name, context.pop('checkpoint_dir'))
+        return context

mmtune/mm/hooks/checkpoint.py

@@ -10,6 +10,7 @@
 from mmcv.runner.dist_utils import master_only
 from mmcv.runner.hooks import CheckpointHook as _CheckpointHook
 from ray.tune.integration.torch import distributed_checkpoint_dir
+from torch.optim import Optimizer
 
 
 @HOOKS.register_module()
@@ -85,8 +86,13 @@ def _save_checkpoint(self, runner: BaseRunner) -> None:
                 The runner to save checkpoints.
         """
         model = runner.model
+        optimizer = runner.optimizer
 
-        meta = dict(mmcv_version=mmcv.__version__, time=time.asctime())
+        meta = dict(
+            mmcv_version=mmcv.__version__,
+            time=time.asctime(),
+            epoch=runner.epoch + 1,
+            iter=runner.iter)
         if is_module_wrapper(model):
             model = model.module
         if hasattr(model, 'CLASSES') and model.CLASSES is not None:
@@ -97,6 +103,13 @@ def _save_checkpoint(self, runner: BaseRunner) -> None:
             'state_dict': weights_to_cpu(get_state_dict(model))
         }
 
+        if isinstance(optimizer, Optimizer):
+            checkpoint['optimizer'] = optimizer.state_dict()
+        elif isinstance(optimizer, dict):
+            checkpoint['optimizer'] = {}
+            for name, optim in optimizer.items():
+                checkpoint['optimizer'][name] = optim.state_dict()
+
         with distributed_checkpoint_dir(
                 step=self.get_iter(runner)) as checkpoint_dir:
             path = os.path.join(checkpoint_dir, 'ray_checkpoint.pth')

tests/test_mm/test_hooks.py

@@ -21,13 +21,15 @@ def test_raycheckpointhook():
     mock_runner = MagicMock()
     mock_runner.inner_iter = 3
     mock_runner.iter = 5
+    mock_runner.epoch = 5
 
     cur_iter = hook.get_iter(mock_runner, False)
     assert cur_iter == 6
     cur_iter = hook.get_iter(mock_runner, True)
     assert cur_iter == 4
 
     mock_runner.model = torch.nn.Linear(2, 2)
+    mock_runner.optimizer = torch.optim.Adam(mock_runner.model.parameters())
 
     hook._save_checkpoint(mock_runner)
     assert os.path.exists('ray_checkpoint.pth')

tests/test_mm/test_rewriters.py

@@ -9,6 +9,7 @@
                                          BaseRewriter, BatchConfigPatcher,
                                          CustomHookRegister, Dump,
                                          InstantiateCfg, MergeConfig,
+                                         ResumeFromCkpt,
                                          SequeunceConfigPatcher)
 from mmtune.mm.context.rewriters.builder import build_rewriter
 from mmtune.utils import dump_cfg
@@ -116,5 +117,15 @@ def test_register():
     cfg = MagicMock()
     cfg.custom_hooks = []
     context = dict(cfg=cfg)
+
     context = register(context)
     assert context['cfg'].custom_hooks == post_custom_hooks
+
+
+def test_resume_ckpt():
+    args = MagicMock()
+    context = dict(args=args, checkpoint_dir='test')
+
+    resume_from_ckpt = ResumeFromCkpt()
+    context = resume_from_ckpt(context)
+    assert context.get('args').resume_from == 'test'