Refactor spaces

Author: KKIEEK

mmtune/ray/spaces/__init__.py

@@ -1,6 +1,10 @@
-from .base import BaseSpace
+from .base import (BaseSpace, Lograndint, Loguniform, Qlograndint, Qloguniform,
+                   Qrandint, Qrandn, Quniform, Randint, Randn, Uniform)
 from .builder import SPACES, build_space
 from .choice import Choice
-from .const import Constant
 
-__all__ = ['BaseSpace', 'SPACES', 'build_space', 'Choice', 'Constant']
+__all__ = [
+    'BaseSpace', 'Uniform', 'Quniform', 'Loguniform', 'Qloguniform', 'Randn',
+    'Qrandn', 'Randint', 'Qrandint', 'Lograndint', 'Qlograndint', 'SPACES',
+    'build_space', 'Choice'
+]

mmtune/ray/spaces/base.py

@@ -1,10 +1,69 @@
 from abc import ABCMeta
+from numbers import Number
+from typing import Callable, Union
+
+import ray.tune as tune
+
+from .builder import SPACES
 
 
 class BaseSpace(metaclass=ABCMeta):
     """Base Space class."""
+    sample: Callable
+
+    def __init__(self, **kwargs):
+        self.kwargs = kwargs
 
     @property
-    def space(self) -> callable:
-        """Return the space."""
-        return getattr(self, '_space', None)
+    def space(self) -> Union[Number, list]:
+        return self.sample.__func__(**self.kwargs)
+
+
+@SPACES.register_module()
+class Uniform(BaseSpace):
+    sample: Callable = tune.uniform
+
+
+@SPACES.register_module()
+class Quniform(BaseSpace):
+    sample: Callable = tune.quniform
+
+
+@SPACES.register_module()
+class Loguniform(BaseSpace):
+    sample: Callable = tune.loguniform
+
+
+@SPACES.register_module()
+class Qloguniform(BaseSpace):
+    sample: Callable = tune.qloguniform
+
+
+@SPACES.register_module()
+class Randn(BaseSpace):
+    sample: Callable = tune.randn
+
+
+@SPACES.register_module()
+class Qrandn(BaseSpace):
+    sample: Callable = tune.qrandn
+
+
+@SPACES.register_module()
+class Randint(BaseSpace):
+    sample: Callable = tune.randint
+
+
+@SPACES.register_module()
+class Qrandint(BaseSpace):
+    sample: Callable = tune.qrandint
+
+
+@SPACES.register_module()
+class Lograndint(BaseSpace):
+    sample: Callable = tune.lograndint
+
+
+@SPACES.register_module()
+class Qlograndint(BaseSpace):
+    sample: Callable = tune.qlograndint

mmtune/ray/spaces/builder.py

@@ -1,43 +1,26 @@
-import inspect
-from typing import Callable, Dict
+from typing import Mapping, Sequence
 
 from mmcv.utils import Registry
-from ray.tune import sample
-
-from .base import BaseSpace
 
 SPACES = Registry('spaces')
 
 
-def _register_space(space: Callable) -> None:
-    """Register a space.
-
-    Args:
-        space (Callable): The space to register.
-    """
-
-    @SPACES.register_module(name=space.__name__.capitalize())
-    class _ImplicitSpace(BaseSpace):
-
-        def __init__(self, *args, **kwargs):
-            self._space = space(*args, **kwargs)
-
-
-for space_name in dir(sample):
-    space = getattr(sample, space_name)
-    if not inspect.isfunction(space):
-        continue
-    _register_space(space)
-
-
-def build_space(cfgs: Dict) -> Dict:
+def build_space(cfg: dict) -> dict:
     """Build a space.
 
     Args:
-        cfgs (Dict): The configurations of the space.
+        cfg (dict): The configurations of the space.
 
     Returns:
-        Dict: The instantiated space.
+        dict: The instantiated space.
     """
-
-    return {key: SPACES.build(cfg).space for key, cfg in cfgs.items()}
+    # inductive step
+    cfg = cfg.copy()
+    for key, val in cfg.items():
+        if isinstance(val, Sequence):
+            cfg[key] = [build_space(v) for v in val if isinstance(v, Mapping)]
+        elif isinstance(val, Mapping):
+            cfg[key] = build_space(val)
+            if cfg.get('type', '') in SPACES:
+                cfg[key] = SPACES.build(cfg).space
+    return cfg

mmtune/ray/spaces/choice.py

@@ -1,14 +1,16 @@
-from typing import Optional, Sequence
+from numbers import Number
+from typing import Callable, Optional, Sequence, Union
 
-from ray.tune.sample import choice
+import ray.tune as tune
 
 from mmtune.utils import ImmutableContainer
 from .base import BaseSpace
 from .builder import SPACES
 
 
-@SPACES.register_module(force=True)
+@SPACES.register_module()
 class Choice(BaseSpace):
+    sample: Callable = tune.choice
 
     def __init__(self,
                  categories: Sequence,
@@ -23,11 +25,16 @@ def __init__(self,
             use_container (bool):
                 Whether to use containers. Defaults to True.
         """
-
         if alias is not None:
             assert len(categories) == len(alias)
-        categories = [
-            ImmutableContainer(c, None if alias is None else alias[idx])
-            if use_container else c for idx, c in enumerate(categories)
-        ]
-        self._space = choice(categories)
+
+        if use_container:
+            aliases = alias or [None] * len(categories)
+            categories = [
+                ImmutableContainer(*it) for it in zip(categories, aliases)
+            ]
+        self.categories = categories
+
+    @property
+    def space(self) -> Union[Number, list]:
+        return self.sample.__func__(self.categories)

mmtune/ray/spaces/const.py

@@ -0,0 +1,24 @@
+from numbers import Number
+from typing import Callable, Union
+
+import ray.tune as tune
+
+from .base import BaseSpace
+from .builder import SPACES
+
+
+@SPACES.register_module()
+class SampleFrom(BaseSpace):
+    sample: Callable = tune.sample_from
+
+    def __init__(self, func: Union[str, Callable], imports=None):
+        if isinstance(func, str):
+            func = eval(func)
+        self.func = func
+        self.imports = imports or []
+
+    @property
+    def space(self) -> Union[Number, list]:
+        for module in self.imports:
+            exec(f'import {module}')
+        return self.sample.__func__(self.func)