diff --git a/package/samplers/hebo_base_sampler/LICENSE b/package/samplers/hebo_base_sampler/LICENSE new file mode 100644 index 00000000..f763c760 --- /dev/null +++ b/package/samplers/hebo_base_sampler/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2024 Hiroki Takizawa + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/package/samplers/hebo_base_sampler/__init__.py b/package/samplers/hebo_base_sampler/__init__.py new file mode 100644 index 00000000..b3a2fabf --- /dev/null +++ b/package/samplers/hebo_base_sampler/__init__.py @@ -0,0 +1,4 @@ +from .sampler import HEBOSampler + + +__all__ = ["HEBOSampler"] diff --git a/package/samplers/hebo_base_sampler/example.py b/package/samplers/hebo_base_sampler/example.py new file mode 100644 index 00000000..96baf41d --- /dev/null +++ b/package/samplers/hebo_base_sampler/example.py @@ -0,0 +1,24 @@ +import optuna +import optunahub +import time + + +module = optunahub.load_module("samplers/hebo_base_sampler") +HEBOSampler = module.HEBOSampler + + +def objective(trial: optuna.trial.Trial) -> float: + x = trial.suggest_float("x", -10, 10) + y = trial.suggest_int("y", -10, 10) + time.sleep(1.0) + return x**2 + y**2 + + +if __name__ == "__main__": + sampler = HEBOSampler(constant_liar=True) + study = optuna.create_study(sampler=sampler) + study.optimize(objective, n_trials=100, n_jobs=2) + print(study.best_trial.params) + + fig = optuna.visualization.plot_optimization_history(study) + fig.write_image("hebo_optimization_history.png") diff --git a/package/samplers/hebo_base_sampler/requirements.txt b/package/samplers/hebo_base_sampler/requirements.txt new file mode 100644 index 00000000..2ba8ec1a --- /dev/null +++ b/package/samplers/hebo_base_sampler/requirements.txt @@ -0,0 +1,3 @@ +optuna +optunahub +hebo@git+https://github.com/huawei-noah/HEBO.git@v0.3.6#subdirectory=HEBO diff --git a/package/samplers/hebo_base_sampler/sampler.py b/package/samplers/hebo_base_sampler/sampler.py new file mode 100644 index 00000000..9169f878 --- /dev/null +++ b/package/samplers/hebo_base_sampler/sampler.py @@ -0,0 +1,146 @@ +from __future__ import annotations + +import numpy as np +import optuna +import pandas as pd +from optuna.distributions import (BaseDistribution, CategoricalDistribution, + FloatDistribution, IntDistribution) +from optuna.samplers import BaseSampler +from optuna.search_space import IntersectionSearchSpace +from optuna.study import Study +from optuna.study._study_direction import StudyDirection +from optuna.trial import FrozenTrial, TrialState + +from hebo.design_space.design_space import DesignSpace +from hebo.optimizers.hebo import HEBO + + +class HEBOSampler(BaseSampler): # type: ignore + def __init__( + self, + seed: int | None = None, + constant_liar: bool = False, + independent_sampler: BaseSampler | None = None, + ) -> None: + self._seed = seed + self._intersection_search_space = IntersectionSearchSpace() + self._independent_sampler = ( + independent_sampler or optuna.samplers.RandomSampler(seed=seed) + ) + self._constant_liar = constant_liar + + def sample_relative( + self, + study: Study, + trial: FrozenTrial, + search_space: dict[str, BaseDistribution], + ) -> dict[str, float]: + if self._constant_liar: + target_states = [TrialState.COMPLETE, TrialState.RUNNING] + else: + target_states = [TrialState.COMPLETE] + trials = study.get_trials(deepcopy=False, states=target_states) + if len(t for t in trials if t.state == TrialState.COMPLETE) < 1: + return {} + + # Assume that the back-end HEBO implementation aims to minimize. + if study.direction == StudyDirection.MINIMIZE: + worst_values = max( + t.values for t in trials if t.state == TrialState.COMPLETE + ) + else: + worst_values = min( + t.values for t in trials if t.state == TrialState.COMPLETE + ) + sign = 1.0 if study.direction == StudyDirection.MINIMIZE else -1.0 + + hebo = HEBO(self._convert_to_hebo_design_space(search_space)) + for t in trials: + hebo_params = {name: t.params[name] for name in search_space.keys()} + if t.state == TrialState.COMPLETE: + hebo.observe(pd.DataFrame([hebo_params]), np.asarray([t.values * sign])) + elif t.state == TrialState.RUNNING: + # If `constant_liar == True`, assume that the RUNNING params result in bad values, + # thus preventing the simultaneous suggestion of (almost) the same params + # during parallel execution. + hebo.observe(pd.DataFrame([hebo_params]), np.asarray([worst_values])) + else: + assert False + params_pd = hebo.suggest() + params = {} + for name in search_space.keys(): + params[name] = params_pd[name].to_numpy()[0] + return params + + def _convert_to_hebo_design_space( + self, search_space: dict[str, BaseDistribution] + ) -> DesignSpace: + design_space = [] + for name, distribution in search_space.items(): + if isinstance(distribution, FloatDistribution) and not distribution.log: + design_space.append( + { + "name": name, + "type": "num", + "lb": distribution.low, + "ub": distribution.high, + } + ) + elif isinstance(distribution, FloatDistribution) and distribution.log: + design_space.append( + { + "name": name, + "type": "pow", + "lb": distribution.low, + "ub": distribution.high, + } + ) + elif isinstance(distribution, IntDistribution) and distribution.log: + design_space.append( + { + "name": name, + "type": "pow_int", + "lb": distribution.low, + "ub": distribution.high, + } + ) + elif isinstance(distribution, IntDistribution) and distribution.step: + design_space.append( + { + "name": name, + "type": "step_int", + "lb": distribution.low, + "ub": distribution.high, + "step": distribution.step, + } + ) + elif isinstance(distribution, IntDistribution): + design_space.append( + { + "name": name, + "type": "int", + "lb": distribution.low, + "ub": distribution.high, + } + ) + elif isinstance(distribution, CategoricalDistribution): + design_space.append( + { + "name": name, + "type": "cat", + "categories": distribution.choices, + } + ) + else: + raise NotImplementedError(f"Unsupported distribution: {distribution}") + return DesignSpace().parse(design_space) + + def infer_relative_search_space(self, study, trial): + return optuna.search_space.intersection_search_space( + study.get_trials(deepcopy=False) + ) + + def sample_independent(self, study, trial, param_name, param_distribution): + return self._independent_sampler.sample_independent( + study, trial, param_name, param_distribution + )