[Nano HPO ] refactor : move optuna related code to deps (#4570)

* refactor to move optuna backends to deps

* remove optuna from docs and refactor to move pruning callback to deps

* fix code style

* fix bug in test

python/nano/src/bigdl/nano/automl/hpo/backend.py

@@ -15,7 +15,6 @@
 #
 
 from enum import Enum
-import optuna
 
 
 class SamplerType(Enum):
@@ -42,97 +41,19 @@ class PrunerType(Enum):
     Threshold = 7
 
 
-class OptunaBackend(object):
-    """A Wrapper to shield user from Optuna specific configurations and API\
-      Later may support other HPO search engines."""
+def create_hpo_backend():
+    """Create HPO Search Backend. Only Optuna is supported."""
+    from bigdl.nano.deps.automl.hpo_api import create_optuna_backend
+    return create_optuna_backend()
 
-    pruner_map = {
-        PrunerType.HyperBand: optuna.pruners.HyperbandPruner,
-        PrunerType.Median: optuna.pruners.MedianPruner,
-        PrunerType.Nop: optuna.pruners.NopPruner,
-        PrunerType.Patient: optuna.pruners.PatientPruner,
-        PrunerType.Percentile: optuna.pruners.PercentilePruner,
-        PrunerType.SuccessiveHalving: optuna.pruners.SuccessiveHalvingPruner,
-        PrunerType.Threshold: optuna.pruners.ThresholdPruner,
-    }
 
-    sampler_map = {
-        SamplerType.TPE: optuna.samplers.TPESampler,
-        SamplerType.CmaEs: optuna.samplers.CmaEsSampler,
-        SamplerType.Grid: optuna.samplers.GridSampler,
-        SamplerType.Random: optuna.samplers.RandomSampler,
-        SamplerType.PartialFixed: optuna.samplers.PartialFixedSampler,
-        SamplerType.NSGAII: optuna.samplers.NSGAIISampler,
-        SamplerType.MOTPE: optuna.samplers.MOTPESampler,
-    }
+def create_tfkeras_pruning_callback(*args, **kwargs):
+    """Create Tensorflow Pruning Callback. Optuna Only."""
+    from bigdl.nano.deps.automl.hpo_api import create_optuna_tfkeras_pruning_callback
+    return create_optuna_tfkeras_pruning_callback()
 
-    @staticmethod
-    def get_other_args(kwargs, kwspaces):
-        """Get key-word arguments which are not search spaces."""
-        return{k: kwargs[k] for k in set(kwargs) - set(kwspaces)}
 
-    @staticmethod
-    def get_hpo_config(trial, configspace):
-        """Get hyper parameter suggestions from search space settings."""
-        # TODO better ways to map ConfigSpace to optuna spaces
-        # fix order of hyperparams in configspace.
-        hp_ordering = configspace.get_hyperparameter_names()
-        config = {}
-        for hp in hp_ordering:
-            hp_obj = configspace.get_hyperparameter(hp)
-            hp_prefix = hp_obj.meta.setdefault('prefix', None)
-            hp_name = hp_prefix + ':' + hp if hp_prefix else hp
-            hp_type = str(type(hp_obj)).lower()  # type of hyperparam
-            if 'integer' in hp_type:
-                hp_dimension = trial.suggest_int(
-                    name=hp_name, low=int(hp_obj.lower), high=int(hp_obj.upper))
-            elif 'float' in hp_type:
-                if hp_obj.log:  # log10-scale hyperparmeter
-                    hp_dimension = trial.suggest_loguniform(
-                        name=hp_name, low=float(hp_obj.lower), high=float(hp_obj.upper))
-                else:
-                    hp_dimension = trial.suggest_float(
-                        name=hp_name, low=float(hp_obj.lower), high=float(hp_obj.upper))
-            elif 'categorical' in hp_type:
-                hp_dimension = trial.suggest_categorical(
-                    name=hp_name, choices=hp_obj.choices)
-            elif 'ordinal' in hp_type:
-                hp_dimension = trial.suggest_categorical(
-                    name=hp_name, choices=hp_obj.sequence)
-            else:
-                raise ValueError("unknown hyperparameter type: %s" % hp)
-            config[hp_name] = hp_dimension
-        return config
-
-    @staticmethod
-    def instantiate(trial, lazyobj):
-        """Instantiate a lazyobject from a trial's sampled param set."""
-        config = OptunaBackend.gen_config(trial, lazyobj)
-        return lazyobj.sample(**config)
-
-    @staticmethod
-    def gen_config(trial, automl_obj):
-        """Generate the param config from a trial's sampled param set."""
-        configspace = automl_obj.cs
-        config = OptunaBackend.get_hpo_config(trial, configspace)
-        other_kwargs = OptunaBackend.get_other_args(
-            automl_obj.kwargs, automl_obj.kwspaces)
-        config.update(other_kwargs)
-        return config
-
-    @staticmethod
-    def create_sampler(sampler_type, kwargs):
-        """Create a hyperparameter sampler by type."""
-        sampler_class = OptunaBackend.sampler_map.get(sampler_type)
-        return sampler_class(**kwargs)
-
-    @staticmethod
-    def create_pruner(pruner_type, kwargs):
-        """Create a pruner by type."""
-        pruner_class = OptunaBackend.pruner_map.get(pruner_type)
-        return pruner_class(**kwargs)
-
-    @staticmethod
-    def create_study(**kwargs):
-        """Create a study to drive the hyperparameter search."""
-        return optuna.create_study(**kwargs)
+def create_pl_pruning_callback(*args, **kwargs):
+    """Create PyTorchLightning Pruning Callback. Optuna Only."""
+    from bigdl.nano.deps.automl.hpo_api import create_optuna_pl_pruning_callback
+    return create_optuna_pl_pruning_callback()

python/nano/src/bigdl/nano/automl/hpo/callgraph.py

@@ -16,7 +16,6 @@
 
 
 from .space import AutoObject
-from .backend import OptunaBackend
 from enum import Enum
 
 
@@ -130,7 +129,7 @@ def _update_cache_from_input(cache, inp):
         return cur_cache
 
     @staticmethod
-    def execute(inputs, outputs, trial):
+    def execute(inputs, outputs, trial, backend):
         """
         Execute the function calls and construct the tensor graph.
 
@@ -173,7 +172,7 @@ def _process_arguments(arguments, cache):
                     arguments, out_cache)
                 # layer is an auto object
                 assert(isinstance(caller, AutoObject))
-                instance = OptunaBackend.instantiate(trial, caller)
+                instance = backend.instantiate(trial, caller)
                 # the actual excution of the functional API
                 out_tensor = instance(new_arguments)
             elif call_type == CALLTYPE.FUNC_SLICE:
@@ -183,15 +182,15 @@ def _process_arguments(arguments, cache):
                 # the actual excution of the functional API
                 out_tensor = source_tensor.__getitem__(*slice_args, **slice_kwargs)
             elif call_type == CALLTYPE.FUNC_CALL:
-                # out_tensor = OptunaBackend.instantiate(trial, caller)
+                # out_tensor = backend.instantiate(trial, caller)
                 new_arguments = _process_arguments(
                     arguments, out_cache)
                 assert(isinstance(caller, AutoObject))
                 # assume tensors does not exist in kwargs
                 # replace only the non-kwargs with new_arguments
                 # TODO revisit to validate the parent tensors in kwargs
                 caller.args, caller.kwargs = new_arguments
-                out_tensor = OptunaBackend.instantiate(trial, caller)
+                out_tensor = backend.instantiate(trial, caller)
             else:
                 raise ValueError("Unexpected CallType: %s" % type)
             out_cache.add_tensor(caller, out_tensor)

python/nano/src/bigdl/nano/automl/hpo/decorator.py

@@ -31,14 +31,14 @@
 import multiprocessing as mp
 import ConfigSpace as CS
 
-from .backend import OptunaBackend
 
 from .space import *
 from .space import _add_hp, _add_cs, _rm_hp, _strip_config_space, SPLITTER
 from .callgraph import CallCache
 
 from bigdl.nano.automl.utils import EasyDict as ezdict
 from bigdl.nano.automl.utils import proxy_methods
+from bigdl.nano.automl.hpo.backend import create_hpo_backend
 
 __all__ = ['args', 'obj', 'func', 'tfmodel', 'plmodel', 'sample_config']
 
@@ -454,7 +454,7 @@ def _model_build(self, trial):
                 # override _model_build to build
                 # the model directly instead of using
                 # modeld_init and model_compile
-                model = OptunaBackend.instantiate(trial, self._lazyobj)
+                model = self.backend.instantiate(trial, self._lazyobj)
                 self._model_compile(model, trial)
                 return model
 
@@ -495,6 +495,7 @@ def __init__(self, **kwargs):
                     elif k in default_config:
                         super_kwargs[k] = default_config[k]
                 super().__init__(**super_kwargs)
+                self.backend = create_hpo_backend()
 
             def __repr__(self):
                 return 'PlAutoMdl -- ' + Cls.__name__
@@ -503,7 +504,7 @@ def _model_build(self, trial):
                 # override _model_build to build
                 # the model directly instead of using
                 # modeld_init and model_compile
-                model = OptunaBackend.instantiate(trial, self._lazyobj)
+                model = self.backend.instantiate(trial, self._lazyobj)
                 # self._model_compile(model, trial)
                 return model
 

python/nano/src/bigdl/nano/automl/hpo/search.py

@@ -30,7 +30,7 @@ def _search_summary(study):
     """
     Print statistics of trials and retrieve the summary for further analysis.
 
-    :param study: the optuna study object
+    :param study: the hpo study object
     :return : the summary object (current we return the study directly, so that
         it allows better flexiblity to do visualization and futher analysis)
     """
@@ -55,7 +55,7 @@ def _end_search(study, model_builder, use_trial_id=-1):
 
     Use the specified trial or best trial to init and compile the base model.
 
-    :param study: the optuna study object.
+    :param study: the hpo study object.
     :param model_builder: the function to build the model.
     :param use_trial_id: int(optional) params of which trial to be used. Defaults to -1.
     :raises ValueError: if study is None.

python/nano/src/bigdl/nano/automl/pytorch/hposearcher.py

@@ -15,14 +15,11 @@
 #
 
 from typing import Any, Dict, Optional, Union
-from numpy import isin
 import pytorch_lightning as pl
 
-
 from pytorch_lightning.trainer.states import TrainerFn, TrainerStatus
 
-
-from bigdl.nano.automl.hpo.backend import OptunaBackend
+from bigdl.nano.automl.hpo.backend import create_hpo_backend
 from .objective import Objective
 from ..hpo.search import (
     _search_summary,
@@ -57,6 +54,7 @@ def __init__(self, trainer: "pl.Trainer") -> None:
         self.study = None
         self.tune_end = False
         self._lazymodel = None
+        self.backend = create_hpo_backend()
 
     def _run_search(self,
                     resume,
@@ -84,20 +82,20 @@ def _run_search(self,
             sampler_type = study_create_kwargs.get('sampler', None)
             if sampler_type:
                 sampler_args = study_create_kwargs.get('sampler_kwargs', {})
-                sampler = OptunaBackend.create_sampler(sampler_type, sampler_args)
+                sampler = self.backend.create_sampler(sampler_type, sampler_args)
                 study_create_kwargs['sampler'] = sampler
                 study_create_kwargs.pop('sampler_kwargs', None)
 
             pruner_type = study_create_kwargs.get('pruner', None)
             if pruner_type:
                 pruner_args = study_create_kwargs.get('pruner_kwargs', {})
-                pruner = OptunaBackend.create_pruner(pruner_type, pruner_args)
+                pruner = self.backend.create_pruner(pruner_type, pruner_args)
                 study_create_kwargs['pruner'] = pruner
                 study_create_kwargs.pop('pruner_kwargs', None)
 
             study_create_kwargs['load_if_exists'] = load_if_exists
             # create study
-            self.study = OptunaBackend.create_study(**study_create_kwargs)
+            self.study = self.backend.create_study(**study_create_kwargs)
 
         # renamed callbacks to tune_callbacks to avoid conflict with fit param
         study_optimize_kwargs = _filter_tuner_args(kwargs, run_keys)
@@ -191,7 +189,7 @@ def search_summary(self):
         """
         Retrive a summary of trials.
 
-        :return: A summary of all the trials. Currently the optuna study is
+        :return: A summary of all the trials. Currently the entire study is
             returned to allow more flexibility for further analysis and visualization.
         """
         return _search_summary(self.study)

python/nano/src/bigdl/nano/automl/pytorch/objective.py

@@ -23,7 +23,7 @@
 from pytorch_lightning.utilities import rank_zero_deprecation
 from pytorch_lightning.core.datamodule import LightningDataModule
 
-from optuna.integration import PyTorchLightningPruningCallback
+from bigdl.nano.automl.hpo.backend import create_pl_pruning_callback
 import inspect
 import copy
 
@@ -33,7 +33,7 @@ def _is_creator(model):
 
 
 class Objective(object):
-    """The Tuning objective for Optuna."""
+    """The Tuning objective for HPO."""
 
     def __init__(self,
                  searcher,
@@ -73,7 +73,7 @@ def _pre_train(self, model, trial):
         # specific args TODO: may need to handle more cases
         if self.pruning:
             callbacks = self.searcher.trainer.callbacks or []
-            pruning_cb = PyTorchLightningPruningCallback(trial, monitor=self.target_metric)
+            pruning_cb = create_pl_pruning_callback(trial, monitor=self.target_metric)
             callbacks.append(pruning_cb)
             self.searcher.trainer.callbacks = callbacks
 
@@ -120,7 +120,7 @@ def __call__(self, trial):
         """
         Execute Training and return target metric in each trial.
 
-        :param: trial: optuna trial which provides the hyperparameter combinition.
+        :param: trial: a trial object which provides the hyperparameter combinition.
         :return: the target metric value.
         """
         # fit

python/nano/src/bigdl/nano/automl/tf/keras/Model.py

@@ -44,7 +44,8 @@ def _model_init_args(self, trial):
         in_tensors, out_tensors = CallCache.execute(
             self.lazyinputs_,
             self.lazyoutputs_,
-            trial)
+            trial,
+            self.backend)
         self.kwargs['inputs'] = in_tensors
         self.kwargs['outputs'] = out_tensors
         return self.kwargs

python/nano/src/bigdl/nano/automl/tf/keras/Sequential.py

@@ -19,7 +19,6 @@
 
 from bigdl.nano.automl.utils import proxy_methods
 from bigdl.nano.automl.tf.mixin import HPOMixin
-from bigdl.nano.automl.hpo.backend import OptunaBackend
 from bigdl.nano.automl.hpo.space import AutoObject
 
 
@@ -59,7 +58,7 @@ def _model_init_args(self, trial):
         instantiated_layers = []
         for layer in self.lazylayers_:
             if isinstance(layer, AutoObject):
-                newl = OptunaBackend.instantiate(trial, layer)
+                newl = self.backend.instantiate(trial, layer)
             else:
                 newl = copy.deepcopy(layer)
             instantiated_layers.append(newl)