[Feature] Support HyperBand and BOHB scheduler (#101)

* Bump ray from 1.9.1 to 2.1.0

Signed-off-by: Junhwa Song <[email protected]>

README.md

@@ -28,11 +28,12 @@ SIATune is an open-source deep learning model hyperparameter tuning toolbox espe
 - **Schedule multiple experiments**
 
   Various scheduling techniques are supported to efficiently manage many experiments.
-  - [x] [AsyncHyperBandScheduler](https://arxiv.org/abs/1810.05934)
-  - [ ] [PopulationBasedTraining](https://www.deepmind.com/blog/population-based-training-of-neural-networks)
-  - [ ] [MedianStoppingRule](https://research.google.com/pubs/pub46180.html)
+  - [x] [Asynchronous HyperBand](https://arxiv.org/abs/1810.05934)
+  - [x] [HyperBand](https://arxiv.org/abs/1603.06560)
+  - [ ] [Median Stopping Rule](https://research.google.com/pubs/pub46180.html)
+  - [ ] [Population Based Training](https://www.deepmind.com/blog/population-based-training-of-neural-networks)
   - [ ] [Population Based Bandits](https://arxiv.org/abs/2002.02518)
-  - [ ] [HyperBandScheduler](https://arxiv.org/abs/1603.06560)
+  - [x] [Bayesian Optimization and HyperBand](https://arxiv.org/abs/1807.01774)
 
 
 - **Distributed tuning system based on Ray**

configs/_base_/scheduler/bohb.py

@@ -0,0 +1,2 @@
+trial_scheduler = dict(
+    type='HyperBandForBOHB', time_attr='training_iteration', max_t=20)

configs/_base_/scheduler/hb.py

@@ -0,0 +1,2 @@
+trial_scheduler = dict(
+    type='HyperBandScheduler', time_attr='training_iteration', max_t=20)

configs/_base_/searcher/bohb.py

@@ -0,0 +1 @@
+searcher = dict(type='TuneBOHB')

configs/_base_/searcher/nevergrad_oneplusone.py

@@ -1,4 +1 @@
-searcher = dict(
-    type='NevergradSearch',
-    budget=256,
-)
+searcher = dict(type='NevergradSearch', budget=256)

configs/_base_/searcher/nevergrad_pso.py

@@ -1,5 +1 @@
-searcher = dict(
-    type='NevergradSearch',
-    optimizer='PSO',
-    budget=256,
-)
+searcher = dict(type='NevergradSearch', optimizer='PSO', budget=256)

requirements/optional.txt

@@ -1,5 +1,7 @@
 bayesian-optimization==1.2.0
+ConfigSpace
 flaml==1.0.14
+hpbandster
 hyperopt==0.2.5
 mlflow==1.23.1
 nevergrad==0.4.3.post7

siatune/ray/schedulers/__init__.py

@@ -1,5 +1,5 @@
 # Copyright (c) SI-Analytics. All rights reserved.
-from .builder import SCHEDULERS, build_scheduler
+from .builder import TRIAL_SCHEDULERS, build_scheduler
 from .pbt import PopulationBasedTraining
 
-__all__ = ['SCHEDULERS', 'build_scheduler', 'PopulationBasedTraining']
+__all__ = ['TRIAL_SCHEDULERS', 'build_scheduler', 'PopulationBasedTraining']

siatune/ray/schedulers/builder.py

@@ -3,15 +3,16 @@
 
 from mmcv.utils import Config, Registry
 from ray import tune
+from ray.tune.schedulers import TrialScheduler
 
-SCHEDULERS = Registry('schedulers')
+TRIAL_SCHEDULERS = Registry('trial scheduler')
 for v in set(tune.schedulers.SCHEDULER_IMPORT.values()):
     if not inspect.isclass(v):
         continue
-    SCHEDULERS.register_module(module=v)
+    TRIAL_SCHEDULERS.register_module(module=v)
 
 
-def build_scheduler(cfg: Config) -> tune.schedulers.TrialScheduler:
+def build_scheduler(cfg: Config) -> TrialScheduler:
     """Build the scheduler from configs.
 
     Args:
@@ -20,4 +21,4 @@ def build_scheduler(cfg: Config) -> tune.schedulers.TrialScheduler:
         tune.schedulers.TrialScheduler: The scheduler.
     """
 
-    return SCHEDULERS.build(cfg)
+    return TRIAL_SCHEDULERS.build(cfg)

siatune/ray/schedulers/pbt.py

@@ -8,7 +8,7 @@
     PopulationBasedTraining as _PopulationBasedTraining
 from ray.tune.search.sample import Domain
 
-from siatune.ray.schedulers import SCHEDULERS
+from siatune.ray.schedulers import TRIAL_SCHEDULERS
 from siatune.ray.spaces import build_space
 from siatune.utils import ImmutableContainer
 
@@ -50,7 +50,7 @@ def explore(
     return new_config
 
 
-@SCHEDULERS.register_module(force=True)
+@TRIAL_SCHEDULERS.register_module(force=True)
 class PopulationBasedTraining(_PopulationBasedTraining):
 
     def __init__(self, *args, **kwargs) -> None:

siatune/ray/searchers/__init__.py

@@ -1,10 +1,5 @@
 # Copyright (c) SI-Analytics. All rights reserved.
 from .builder import SEARCHERS, build_searcher
-from .flaml import BlendSearch, CFOSearch
-from .hyperopt import HyperOptSearch
 from .nevergrad import NevergradSearch
 
-__all__ = [
-    'SEARCHERS', 'build_searcher', 'BlendSearch', 'CFOSearch',
-    'HyperOptSearch', 'NevergradSearch'
-]
+__all__ = ['SEARCHERS', 'build_searcher', 'NevergradSearch']

siatune/ray/searchers/builder.py

@@ -1,8 +1,12 @@
 # Copyright (c) SI-Analytics. All rights reserved.
+
 from mmcv.utils import Config, Registry
+from ray import tune
 from ray.tune.search import Searcher
 
 SEARCHERS = Registry('searchers')
+for func in set(tune.search.SEARCH_ALG_IMPORT.values()):
+    SEARCHERS.register_module(module=func())
 
 
 def build_searcher(cfg: Config) -> Searcher:

siatune/ray/searchers/nevergrad.py

@@ -21,7 +21,7 @@
     optimizer_registry = dict()
 
 
-@SEARCHERS.register_module()
+@SEARCHERS.register_module(force=True)
 class NevergradSearch(_NevergradSearch):
     """Search with Nevergrad."""
 
@@ -104,8 +104,8 @@ def _setup_nevergrad(self) -> None:
                 if len(self._nevergrad_opt.instrumentation.args) != 1:
                     raise ValueError(
                         'Instrumented optimizers should use kwargs only')
-        if self._parameters is not None and \
-           self._nevergrad_opt.dimension != len(self._parameters):
+        if self._parameters is not None and (self._nevergrad_opt.dimension !=
+                                             len(self._parameters)):
             raise ValueError('len(parameters_names) must match optimizer '
                              'dimension for non-instrumented optimizers')
 

tests/test_ray/test_schedulers.py

@@ -1,9 +1,9 @@
-from siatune.ray.schedulers import SCHEDULERS, build_scheduler
+from siatune.ray.schedulers import TRIAL_SCHEDULERS, build_scheduler
 
 
 def test_build_schedulers():
 
-    @SCHEDULERS.register_module()
+    @TRIAL_SCHEDULERS.register_module()
     class TestScheduler:
         pass
 

tests/test_ray/test_searchers.py

@@ -1,9 +1,7 @@
 import pytest
 from ray import tune
 
-from siatune.ray.searchers import (SEARCHERS, BlendSearch, CFOSearch,
-                                   HyperOptSearch, NevergradSearch,
-                                   build_searcher)
+from siatune.ray.searchers import SEARCHERS, build_searcher
 
 
 def test_build_searcher():
@@ -39,7 +37,17 @@ def test_blend(trainable, config):
         trainable,
         metric='mean_loss',
         mode='min',
-        search_alg=BlendSearch(),
+        search_alg=build_searcher(dict(type='BlendSearch')),
+        num_samples=2,
+        config=config)
+
+
+def test_bohb(trainable, config):
+    tune.run(
+        trainable,
+        metric='mean_loss',
+        mode='min',
+        search_alg=build_searcher(dict(type='TuneBOHB')),
         num_samples=2,
         config=config)
 
@@ -49,7 +57,7 @@ def test_cfo(trainable, config):
         trainable,
         metric='mean_loss',
         mode='min',
-        search_alg=CFOSearch(),
+        search_alg=build_searcher(dict(type='CFO')),
         num_samples=2,
         config=config)
 
@@ -59,7 +67,7 @@ def test_hyperopt(trainable, config):
         trainable,
         metric='mean_loss',
         mode='min',
-        search_alg=HyperOptSearch(),
+        search_alg=build_searcher(dict(type='HyperOptSearch')),
         num_samples=2,
         config=config)
 
@@ -69,6 +77,15 @@ def test_nevergrad(trainable, config):
         trainable,
         metric='mean_loss',
         mode='min',
-        search_alg=NevergradSearch(optimizer='PSO', budget=2),
+        search_alg=build_searcher(dict(type='NevergradSearch', budget=1)),
+        num_samples=2,
+        config=config)
+
+    tune.run(
+        trainable,
+        metric='mean_loss',
+        mode='min',
+        search_alg=build_searcher(
+            dict(type='NevergradSearch', optimizer='PSO', budget=1)),
         num_samples=2,
         config=config)