[Enhance] Support MIM trainable

configs/_base_/context/train.py

@@ -1,5 +1,5 @@
 task = dict(rewriters=[
-    dict(type='InstantiateCfg', arg_name='config', key='base_cfg'),
+    dict(type='InstantiateCfg', key='base_cfg'),
     dict(type='BatchConfigPatcher', key='searched_cfg'),
     dict(type='SequeunceConfigPatcher', key='searched_cfg'),
     dict(
@@ -12,10 +12,11 @@
         key='cfg',
         post_custom_hooks=[
             dict(
-                type='RayTuneLoggerHook',
+                type='RayTuneReporterHook',
                 filtering_key='val',
                 priority='VERY_LOW'),
-            dict(type='RayCheckpointHook', by_epoch=True, interval=1)
         ]),
-    dict(type='Dump', key='cfg', arg_name='config')
+    dict(type='ResumeFromCkpt'),
+    dict(type='Dump', key='cfg'),
+    dict(type='AttachTrialInfoToPath')
 ])

configs/mmcls/mim_cifar_100_asynchb_nevergrad_pso.py

@@ -0,0 +1,20 @@
+_base_ = [
+    '../_base_/context/train.py', '../_base_/searcher/nevergrad_pso.py',
+    '../_base_/scheduler/asynchb.py', '../_base_/space/mmcls_model.py',
+    '../_base_/space/optimizer.py', '../_base_/space/batch_size.py'
+]
+
+space = {
+    'model': {{_base_.model}},
+    'model.head.num_classes': 100,
+    'optimizer': {{_base_.optimizer}},
+    'data.samples_per_gpu': {{_base_.batch_size}},
+}
+
+task = dict(type='MIM', pkg_name='mmcls')
+tune_cfg = dict(
+    num_samples=8,
+    metric='val/accuracy_top-1',
+    mode='max',
+    reuse_actors=False,
+    chdir_to_trial_dir=False)

siatune/codebase/__init__.py

@@ -4,22 +4,15 @@
 from .builder import TASKS, build_task
 from .cont_test_func import ContinuousTestFunction
 from .disc_test_func import DiscreteTestFunction
+from .mim import MIM
 from .mm import MMBaseTask
 from .mmcls import MMClassification
 from .mmdet import MMDetection
 from .mmedit import MMEditing
 from .mmseg import MMSegmentation
 
 __all__ = [
-    'TASKS',
-    'build_task',
-    'BaseTask',
-    'BlackBoxTask',
-    'ContinuousTestFunction',
-    'DiscreteTestFunction',
-    'MMBaseTask',
-    'MMClassification',
-    'MMDetection',
-    'MMEditing',
-    'MMSegmentation',
+    'TASKS', 'build_task', 'BaseTask', 'BlackBoxTask',
+    'ContinuousTestFunction', 'DiscreteTestFunction', 'MMBaseTask',
+    'MMClassification', 'MMDetection', 'MMEditing', 'MMSegmentation', 'MIM'
 ]

siatune/codebase/base.py

@@ -2,7 +2,7 @@
 import argparse
 from abc import ABCMeta, abstractmethod
 from copy import deepcopy
-from typing import Callable, Optional, Sequence, Union
+from typing import Optional, Sequence, Union
 
 from ray.tune import Trainable
 
@@ -26,7 +26,8 @@ class BaseTask(metaclass=ABCMeta):
         Aggregate the information we define as context,
         convert it into a refined argparse namespace, and input it to run.
         The context consists of:
-            1. args (argparse.Namespace): The low level CLI arguments.
+            1. args (argparse.Namespace | Sequence[str]):
+                The low level CLI arguments.
             2. searched_cfg (Dict):
                 The configuration searched by the algorithm.
         Inputs: searched_cfg (Dict)
@@ -46,11 +47,15 @@ class BaseTask(metaclass=ABCMeta):
 
     def __init__(self,
                  args: Sequence[str],
-                 num_workers: int,
+                 num_workers: int = 1,
                  num_cpus_per_worker: int = 1,
                  num_gpus_per_worker: int = 1,
-                 rewriters: Optional[Union[list, dict]] = None):
-        self.args = self.parse_args(args)
+                 rewriters: Optional[Union[list, dict]] = None,
+                 should_parse: bool = True):
+
+        if should_parse:
+            args = self.parse_args(args)
+        self.args = args
 
         self.num_workers = num_workers
         self.num_cpus_per_worker = num_cpus_per_worker
@@ -61,7 +66,8 @@ def __init__(self,
         self.rewriters = rewriters
 
     @abstractmethod
-    def parse_args(self, args: Sequence[str]) -> argparse.Namespace:
+    def parse_args(self,
+                   args: Sequence[str]) -> Union[argparse.Namespace, None]:
         """Define and parse the necessary arguments for the task.
 
         Args:
@@ -72,15 +78,12 @@ def parse_args(self, args: Sequence[str]) -> argparse.Namespace:
         """
         pass
 
-    def context_aware_run(self, searched_cfg: dict) -> Callable:
+    def context_aware_run(self, searched_cfg: dict):
         """Gather and refine the information received by users and Ray.tune to
         execute the objective task.
 
         Args:
             searched_cfg (Dict): The searched configuration.
-
-        Returns:
-            Callable: The result of the objective task.
         """
 
         context_manager = ContextManager(self.rewriters)

siatune/codebase/mim.py

@@ -0,0 +1,68 @@
+# Copyright (c) SI-Analytics. All rights reserved.
+from importlib.machinery import SourceFileLoader
+from typing import Sequence
+
+from siatune.utils import get_train_script
+from .builder import TASKS
+from .mm import MMBaseTask
+
+
+class _EntrypointExecutor:
+    """Execute the entrypoint of open mm train-based projects.
+
+    Args:
+        pkg_name (str): The abbreviation of the package.
+        argv (Sequence[str]): The arguments for `tools/train.py`
+        module_name (str):
+            The name of the module to execute. Defaults to 'main'.
+    """
+
+    def __init__(self,
+                 pkg_name: str,
+                 argv: Sequence[str],
+                 module_name: str = 'main'):
+        self._train_script = get_train_script(pkg_name)
+        self._module_name = module_name
+        self._argv = argv
+        self._entrypoint = SourceFileLoader(self._module_name,
+                                            self._train_script).load_module()
+
+    def _hijack_argv(self, argv: Sequence[str]):
+        """Hijack the command line arguments.
+
+        Args:
+            argv (Sequence[str]): The arguments for `tools/train.py`
+        """
+        import sys
+        sys.argv[1:] = argv
+        return
+
+    def execute(self):
+        """Run the task."""
+        self._hijack_argv(self._argv)
+        getattr(self._entrypoint, self._module_name)()
+
+
+@TASKS.register_module()
+class MIM(MMBaseTask):
+    """Wrapper class execute any script provided by all OpenMMLab codebases.
+
+    Args:
+        pkg_name (str): The abbreviation of the package.
+    """
+
+    def __init__(self, pkg_name: str, **kwargs):
+        self._pkg_name = pkg_name
+        super().__init__(should_parse=False, **kwargs)
+
+    def parse_args(self, *args, **kwargs) -> None:
+        pass
+
+    def run(self, args: Sequence[str]):
+        """This method runs a task in the MIM framework.
+
+        Args:
+            args (Sequence[str]): A list of command-line arguments.
+        """
+        executor = _EntrypointExecutor(self._pkg_name, args)
+        executor.execute()

siatune/codebase/mm.py

@@ -1,11 +1,12 @@
 # Copyright (c) SI-Analytics. All rights reserved.
 from abc import ABCMeta
+from copy import deepcopy
+from typing import Callable
 
-import torch
-from ray.air.config import ScalingConfig
-from ray.train.data_parallel_trainer import DataParallelTrainer
+from ray.tune import with_resources
 
-from siatune.tune import MMBackendConfig
+from siatune.core import ContextManager, DistributedTorchLauncher
+from siatune.utils import ImmutableContainer
 from .base import BaseTask
 from .builder import TASKS
 
@@ -14,17 +15,36 @@
 class MMBaseTask(BaseTask, metaclass=ABCMeta):
     """Wrap the apis of open mm train-based projects."""
 
-    def create_trainable(self) -> DataParallelTrainer:
-        """Get a :class:`DataParallelTrainer` instance.
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        assert self.num_gpus_per_worker == 1
+        self.launcher = DistributedTorchLauncher(
+            self.num_cpus_per_worker,
+            self.num_workers,
+        )
+
+    def create_trainable(self) -> Callable:
+        """Get a trainable task.
 
         Returns:
-            DataParallelTrainer: Trainer to optimize hyperparameter.
+            Callable: Callable object to optimize hyperparameter.
         """
+        return with_resources(self.context_aware_run, self.launcher.resources)
+
+    def context_aware_run(self, searched_cfg: dict):
+        """Gather and refine the information received by users and Ray.tune to
+        execute the objective task.
+
+        Args:
+            searched_cfg (Dict): The searched configuration.
+        """
+
+        context_manager = ContextManager(self.rewriters)
+        context = dict(
+            args=deepcopy(self.args),
+            searched_cfg=deepcopy(ImmutableContainer.decouple(searched_cfg)),
+        )
+        return context_manager(self.dist_run)(**context)
 
-        return DataParallelTrainer(
-            self.context_aware_run,
-            backend_config=MMBackendConfig(),
-            scaling_config=ScalingConfig(
-                trainer_resources=dict(CPU=self.num_cpus_per_worker),
-                num_workers=self.num_workers,
-                use_gpu=torch.cuda.is_available()))
+    def dist_run(self, *args, **kwargs):
+        self.launcher.launch(self.run, *args, **kwargs)

siatune/core/__init__.py

@@ -1,6 +1,12 @@
 # Copyright (c) SI-Analytics. All rights reserved.
 from .context import ContextManager
 from .hooks import *  # noqa F403
+from .launch import DistributedTorchLauncher
 from .rewriters import REWRITERS, build_rewriter
 
-__all__ = ['ContextManager', 'REWRITERS', 'build_rewriter']
+__all__ = [
+    'ContextManager',
+    'REWRITERS',
+    'build_rewriter',
+    'DistributedTorchLauncher',
+]

siatune/core/hooks/__init__.py

@@ -1,5 +1,4 @@
 # Copyright (c) SI-Analytics. All rights reserved.
-from .checkpoint import RayCheckpointHook
-from .reporter import RayTuneLoggerHook
+from .reporter import RayTuneReporterHook
 
-__all__ = ['RayCheckpointHook', 'RayTuneLoggerHook']
+__all__ = ['RayTuneReporterHook']