From 16486a8df3c25089e4ad8a95abfd37fc3bf9f87f Mon Sep 17 00:00:00 2001 From: krfricke Date: Thu, 13 Aug 2020 01:13:22 +0200 Subject: [PATCH] [tune] Add OptunaSearcher wrapper around Optuna samplers (#10044) Co-authored-by: Richard Liaw Co-authored-by: Kai Fricke --- doc/source/tune/api_docs/suggestion.rst | 13 ++ doc/source/tune/examples/optuna_example.rst | 6 + python/ray/tune/BUILD | 9 ++ python/ray/tune/examples/optuna_example.py | 53 +++++++ python/ray/tune/suggest/optuna.py | 149 ++++++++++++++++++++ python/ray/tune/tests/test_tune_restore.py | 19 ++- python/requirements_tune.txt | 1 + 7 files changed, 249 insertions(+), 1 deletion(-) create mode 100644 doc/source/tune/examples/optuna_example.rst create mode 100644 python/ray/tune/examples/optuna_example.py create mode 100644 python/ray/tune/suggest/optuna.py diff --git a/doc/source/tune/api_docs/suggestion.rst b/doc/source/tune/api_docs/suggestion.rst index 41da1d789..ea2bbfa8c 100644 --- a/doc/source/tune/api_docs/suggestion.rst +++ b/doc/source/tune/api_docs/suggestion.rst @@ -50,6 +50,10 @@ Summary - Gradient-free Optimization - [`Nevergrad `__] - :doc:`/tune/examples/nevergrad_example` + * - :ref:`OptunaSearch ` + - Optuna search algorithms + - [`Optuna `__] + - :doc:`/tune/examples/optuna_example` * - :ref:`ZOOptSearch ` - Zeroth-order Optimization - [`ZOOpt `__] @@ -184,6 +188,15 @@ Nevergrad (tune.suggest.nevergrad.NevergradSearch) .. _`Nevergrad README's Optimization section`: https://github.com/facebookresearch/nevergrad/blob/master/docs/optimization.rst#choosing-an-optimizer +.. _tune-optuna: + +Optuna (tune.suggest.optuna.OptunaSearch) +-------------------------------------------------- + +.. autoclass:: ray.tune.suggest.optuna.OptunaSearch + +.. _`Optuna samplers`: https://optuna.readthedocs.io/en/stable/reference/samplers.html + .. _sigopt: SigOpt (tune.suggest.sigopt.SigOptSearch) diff --git a/doc/source/tune/examples/optuna_example.rst b/doc/source/tune/examples/optuna_example.rst new file mode 100644 index 000000000..296bebb65 --- /dev/null +++ b/doc/source/tune/examples/optuna_example.rst @@ -0,0 +1,6 @@ +:orphan: + +optuna_example +~~~~~~~~~~~~~~~~ + +.. literalinclude:: /../../python/ray/tune/examples/optuna_example.py \ No newline at end of file diff --git a/python/ray/tune/BUILD b/python/ray/tune/BUILD index e786c0448..6cbb391a9 100644 --- a/python/ray/tune/BUILD +++ b/python/ray/tune/BUILD @@ -428,6 +428,15 @@ py_test( args = ["--smoke-test"] ) +py_test( + name = "optuna_example", + size = "medium", + srcs = ["examples/optuna_example.py"], + deps = [":tune_lib"], + tags = ["exclusive", "example"], + args = ["--smoke-test"] +) + py_test( name = "pbt_convnet_example", size = "medium", diff --git a/python/ray/tune/examples/optuna_example.py b/python/ray/tune/examples/optuna_example.py new file mode 100644 index 000000000..6bafbfeb9 --- /dev/null +++ b/python/ray/tune/examples/optuna_example.py @@ -0,0 +1,53 @@ +"""This test checks that Optuna is functional. + +It also checks that it is usable with a separate scheduler. +""" +import time + +import ray +from ray import tune +from ray.tune.schedulers import AsyncHyperBandScheduler +from ray.tune.suggest.optuna import OptunaSearch, param + + +def evaluation_fn(step, width, height): + return (0.1 + width * step / 100)**(-1) + height * 0.1 + + +def easy_objective(config): + # Hyperparameters + width, height = config["width"], config["height"] + + for step in range(config["steps"]): + # Iterative training function - can be any arbitrary training procedure + intermediate_score = evaluation_fn(step, width, height) + # Feed the score back back to Tune. + tune.report(iterations=step, mean_loss=intermediate_score) + time.sleep(0.1) + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument( + "--smoke-test", action="store_true", help="Finish quickly for testing") + args, _ = parser.parse_known_args() + ray.init(configure_logging=False) + + space = [ + param.suggest_uniform("width", 0, 20), + param.suggest_uniform("height", -100, 100), + # This is an ignored parameter. + param.suggest_categorical("activation", ["relu", "tanh"]) + ] + + config = { + "num_samples": 10 if args.smoke_test else 100, + "config": { + "steps": 100, + } + } + algo = OptunaSearch(space, metric="mean_loss", mode="min") + scheduler = AsyncHyperBandScheduler(metric="mean_loss", mode="min") + tune.run(easy_objective, search_alg=algo, scheduler=scheduler, **config) diff --git a/python/ray/tune/suggest/optuna.py b/python/ray/tune/suggest/optuna.py new file mode 100644 index 000000000..51d3efc14 --- /dev/null +++ b/python/ray/tune/suggest/optuna.py @@ -0,0 +1,149 @@ +import logging +import pickle + +from ray.tune.result import TRAINING_ITERATION + +try: + import optuna as ot +except ImportError: + ot = None + +from ray.tune.suggest import Searcher + +logger = logging.getLogger(__name__) + + +class _Param: + def __getattr__(self, item): + def _inner(*args, **kwargs): + return (item, args, kwargs) + + return _inner + + +param = _Param() + + +class OptunaSearch(Searcher): + """A wrapper around Optuna to provide trial suggestions. + + `Optuna `_ is a hyperparameter optimization library. + In contrast to other libraries, it employs define-by-run style + hyperparameter definitions. + + This Searcher is a thin wrapper around Optuna's search algorithms. + You can pass any Optuna sampler, which will be used to generate + hyperparameter suggestions. + + Please note that this wrapper does not support define-by-run, so the + search space will be configured before running the optimization. You will + also need to use a Tune trainable (e.g. using the function API) with + this wrapper. + + For defining the search space, use ``ray.tune.suggest.optuna.param`` + (see example). + + Args: + space (list): Hyperparameter search space definition for Optuna's + sampler. This is a list, and samples for the parameters will + be obtained in order. + metric (str): Metric that is reported back to Optuna on trial + completion. + mode (str): One of {min, max}. Determines whether objective is + minimizing or maximizing the metric attribute. + sampler (optuna.samplers.BaseSampler): Optuna sampler used to + draw hyperparameter configurations. Defaults to ``TPESampler``. + + Example: + + .. code-block:: python + + from ray.tune.suggest.optuna import OptunaSearch, param + + space = [ + param.suggest_uniform("a", 6, 8), + param.suggest_uniform("b", 10, 20) + ] + + algo = OptunaSearch( + space, + metric="loss", + mode="min") + + .. versionadded:: 0.8.8 + + """ + + def __init__( + self, + space, + metric="episode_reward_mean", + mode="max", + sampler=None, + ): + assert ot is not None, ( + "Optuna must be installed! Run `pip install optuna`.") + super(OptunaSearch, self).__init__( + metric=metric, + mode=mode, + max_concurrent=None, + use_early_stopped_trials=None) + + self._space = space + + self._study_name = "optuna" # Fixed study name for in-memory storage + self._sampler = sampler or ot.samplers.TPESampler() + assert isinstance(self._sampler, ot.samplers.BaseSampler), \ + "You can only pass an instance of `optuna.samplers.BaseSampler` " \ + "as a sampler to `OptunaSearcher`." + + self._pruner = ot.pruners.NopPruner() + self._storage = ot.storages.InMemoryStorage() + + self._ot_trials = {} + self._ot_study = ot.study.create_study( + storage=self._storage, + sampler=self._sampler, + pruner=self._pruner, + study_name=self._study_name, + direction="minimize" if mode == "min" else "maximize", + load_if_exists=True) + + def suggest(self, trial_id): + if trial_id not in self._ot_trials: + ot_trial_id = self._storage.create_new_trial( + self._ot_study._study_id) + self._ot_trials[trial_id] = ot.trial.Trial(self._ot_study, + ot_trial_id) + ot_trial = self._ot_trials[trial_id] + params = {} + for (fn, args, kwargs) in self._space: + param_name = args[0] if len(args) > 0 else kwargs["name"] + params[param_name] = getattr(ot_trial, fn)(*args, **kwargs) + return params + + def on_trial_result(self, trial_id, result): + metric = result[self.metric] + step = result[TRAINING_ITERATION] + ot_trial = self._ot_trials[trial_id] + ot_trial.report(metric, step) + + def on_trial_complete(self, trial_id, result=None, error=False): + ot_trial = self._ot_trials[trial_id] + ot_trial_id = ot_trial._trial_id + self._storage.set_trial_value(ot_trial_id, result.get( + self.metric, None)) + self._storage.set_trial_state(ot_trial_id, + ot.trial.TrialState.COMPLETE) + + def save(self, checkpoint_path): + save_object = (self._storage, self._pruner, self._sampler, + self._ot_trials, self._ot_study) + with open(checkpoint_path, "wb") as outputFile: + pickle.dump(save_object, outputFile) + + def restore(self, checkpoint_path): + with open(checkpoint_path, "rb") as inputFile: + save_object = pickle.load(inputFile) + self._storage, self._pruner, self._sampler, \ + self._ot_trials, self._ot_study = save_object diff --git a/python/ray/tune/tests/test_tune_restore.py b/python/ray/tune/tests/test_tune_restore.py index 0286be279..89077da90 100644 --- a/python/ray/tune/tests/test_tune_restore.py +++ b/python/ray/tune/tests/test_tune_restore.py @@ -19,6 +19,7 @@ from ray.tune.suggest.dragonfly import DragonflySearch from ray.tune.suggest.bayesopt import BayesOptSearch from ray.tune.suggest.skopt import SkOptSearch from ray.tune.suggest.nevergrad import NevergradSearch +from ray.tune.suggest.optuna import OptunaSearch, param as ot_param from ray.tune.suggest.sigopt import SigOptSearch from ray.tune.suggest.zoopt import ZOOptSearch from ray.tune.utils import validate_save_restore @@ -292,7 +293,23 @@ class NevergradWarmStartTest(AbstractWarmStartTest, unittest.TestCase): return search_alg, cost -class DragonflyWarmSTartTest(AbstractWarmStartTest, unittest.TestCase): +class OptunaWarmStartTest(AbstractWarmStartTest, unittest.TestCase): + def set_basic_conf(self): + from optuna.samplers import TPESampler + space = [ + ot_param.suggest_uniform("width", 0, 20), + ot_param.suggest_uniform("height", -100, 100) + ] + + def cost(space, reporter): + reporter(loss=(space["height"] - 14)**2 - abs(space["width"] - 3)) + + search_alg = OptunaSearch( + space, sampler=TPESampler(seed=10), metric="loss", mode="min") + return search_alg, cost + + +class DragonflyWarmStartTest(AbstractWarmStartTest, unittest.TestCase): def set_basic_conf(self): from dragonfly.opt.gp_bandit import EuclideanGPBandit from dragonfly.exd.experiment_caller import EuclideanFunctionCaller diff --git a/python/requirements_tune.txt b/python/requirements_tune.txt index 4aaaab323..976149566 100644 --- a/python/requirements_tune.txt +++ b/python/requirements_tune.txt @@ -14,6 +14,7 @@ matplotlib mlflow mxnet nevergrad +optuna pytest-remotedata>=0.3.1 pytorch-lightning scikit-optimize