732 lines
26 KiB
Python
732 lines
26 KiB
Python
import functools
|
|
import inspect
|
|
import logging
|
|
import pickle
|
|
import time
|
|
import warnings
|
|
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
|
|
|
|
from packaging import version
|
|
|
|
from ray.air.constants import TRAINING_ITERATION
|
|
from ray.tune.result import DEFAULT_METRIC
|
|
from ray.tune.search import (
|
|
UNDEFINED_METRIC_MODE,
|
|
UNDEFINED_SEARCH_SPACE,
|
|
UNRESOLVED_SEARCH_SPACE,
|
|
Searcher,
|
|
)
|
|
from ray.tune.search.sample import (
|
|
Categorical,
|
|
Domain,
|
|
Float,
|
|
Integer,
|
|
LogUniform,
|
|
Quantized,
|
|
Uniform,
|
|
)
|
|
from ray.tune.search.variant_generator import parse_spec_vars
|
|
from ray.tune.utils.util import flatten_dict, unflatten_dict, validate_warmstart
|
|
|
|
try:
|
|
import optuna as ot
|
|
from optuna.distributions import BaseDistribution as OptunaDistribution
|
|
from optuna.samplers import BaseSampler
|
|
from optuna.storages import BaseStorage
|
|
from optuna.trial import Trial as OptunaTrial, TrialState as OptunaTrialState
|
|
except ImportError:
|
|
ot = None
|
|
OptunaDistribution = None
|
|
BaseSampler = None
|
|
BaseStorage = None
|
|
OptunaTrialState = None
|
|
OptunaTrial = None
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# print a warning if define by run function takes longer than this to execute
|
|
DEFINE_BY_RUN_WARN_THRESHOLD_S = 1 # 1 is arbitrary
|
|
|
|
|
|
class _OptunaTrialSuggestCaptor:
|
|
"""Utility to capture returned values from Optuna's suggest_ methods.
|
|
|
|
This will wrap around the ``optuna.Trial` object and decorate all
|
|
`suggest_` callables with a function capturing the returned value,
|
|
which will be saved in the ``captured_values`` dict.
|
|
"""
|
|
|
|
def __init__(self, ot_trial: OptunaTrial) -> None:
|
|
self.ot_trial = ot_trial
|
|
self.captured_values: Dict[str, Any] = {}
|
|
|
|
def _get_wrapper(self, func: Callable) -> Callable:
|
|
sig = inspect.signature(func)
|
|
|
|
@functools.wraps(func)
|
|
def wrapper(*args, **kwargs):
|
|
# name is always the first arg for suggest_ methods
|
|
bound = sig.bind_partial(*args, **kwargs)
|
|
bound.apply_defaults()
|
|
if "name" not in bound.arguments:
|
|
raise ValueError("missing required argument: name")
|
|
name = bound.arguments["name"]
|
|
ret = func(*args, **kwargs)
|
|
self.captured_values[name] = ret
|
|
return ret
|
|
|
|
return wrapper
|
|
|
|
def __getattr__(self, item_name: str) -> Any:
|
|
item = getattr(self.ot_trial, item_name)
|
|
if item_name.startswith("suggest_") and callable(item):
|
|
return self._get_wrapper(item)
|
|
return item
|
|
|
|
|
|
class OptunaSearch(Searcher):
|
|
"""A wrapper around Optuna to provide trial suggestions.
|
|
|
|
`Optuna <https://optuna.org/>`_ 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.
|
|
|
|
Multi-objective optimization is supported.
|
|
|
|
.. note::
|
|
``OptunaSearch`` requires ``optuna>=3.0``.
|
|
|
|
Args:
|
|
space: Hyperparameter search space definition for
|
|
Optuna's sampler. This can be either a :class:`dict` with
|
|
parameter names as keys and ``optuna.distributions`` as values,
|
|
or a Callable - in which case, it should be a define-by-run
|
|
function using ``optuna.trial`` to obtain the hyperparameter
|
|
values. The function should return either a :class:`dict` of
|
|
constant values with names as keys, or None.
|
|
For more information, see https://optuna.readthedocs.io\
|
|
/en/stable/tutorial/10_key_features/002_configurations.html.
|
|
|
|
.. warning::
|
|
No actual computation should take place in the define-by-run
|
|
function. Instead, put the training logic inside the function
|
|
or class trainable passed to ``tune.Tuner()``.
|
|
|
|
metric: The training result objective value attribute. If
|
|
None but a mode was passed, the anonymous metric ``_metric``
|
|
will be used per default. Can be a list of metrics for
|
|
multi-objective optimization.
|
|
mode: One of {min, max}. Determines whether objective is
|
|
minimizing or maximizing the metric attribute. Can be a list of
|
|
modes for multi-objective optimization (corresponding to
|
|
``metric``).
|
|
points_to_evaluate: Initial parameter suggestions to be run
|
|
first. This is for when you already have some good parameters
|
|
you want to run first to help the algorithm make better suggestions
|
|
for future parameters. Needs to be a list of dicts containing the
|
|
configurations.
|
|
sampler: Optuna sampler used to
|
|
draw hyperparameter configurations. Defaults to ``TPESampler``,
|
|
which supports both single- and multi-objective optimization.
|
|
See https://optuna.readthedocs.io/en/stable/reference/samplers/index.html
|
|
for available Optuna samplers.
|
|
study_name: Optuna study name that uniquely identifies the trial
|
|
results. Defaults to ``"optuna"``.
|
|
storage: Optuna storage used for storing trial results to
|
|
storages other than in-memory storage,
|
|
for instance optuna.storages.RDBStorage.
|
|
seed: Seed to initialize sampler with. This parameter is only
|
|
used when ``sampler=None``. In all other cases, the sampler
|
|
you pass should be initialized with the seed already.
|
|
evaluated_rewards: If you have previously evaluated the
|
|
parameters passed in as points_to_evaluate you can avoid
|
|
re-running those trials by passing in the reward attributes
|
|
as a list so the optimiser can be told the results without
|
|
needing to re-compute the trial. Must be the same length as
|
|
points_to_evaluate.
|
|
|
|
.. warning::
|
|
When using ``evaluated_rewards``, the search space ``space``
|
|
must be provided as a :class:`dict` with parameter names as
|
|
keys and ``optuna.distributions`` instances as values. The
|
|
define-by-run search space definition is not yet supported with
|
|
this functionality.
|
|
|
|
Tune automatically converts search spaces to Optuna's format:
|
|
|
|
.. code-block:: python
|
|
|
|
from ray.tune.search.optuna import OptunaSearch
|
|
|
|
config = {
|
|
"a": tune.uniform(6, 8)
|
|
"b": tune.loguniform(1e-4, 1e-2)
|
|
}
|
|
|
|
optuna_search = OptunaSearch(
|
|
metric="loss",
|
|
mode="min")
|
|
|
|
tuner = tune.Tuner(
|
|
trainable,
|
|
tune_config=tune.TuneConfig(
|
|
search_alg=optuna_search,
|
|
),
|
|
param_space=config,
|
|
)
|
|
tuner.fit()
|
|
|
|
If you would like to pass the search space manually, the code would
|
|
look like this:
|
|
|
|
.. code-block:: python
|
|
|
|
from ray.tune.search.optuna import OptunaSearch
|
|
import optuna
|
|
|
|
space = {
|
|
"a": optuna.distributions.FloatDistribution(6, 8),
|
|
"b": optuna.distributions.FloatDistribution(1e-4, 1e-2, log=True),
|
|
}
|
|
|
|
optuna_search = OptunaSearch(
|
|
space,
|
|
metric="loss",
|
|
mode="min")
|
|
|
|
tuner = tune.Tuner(
|
|
trainable,
|
|
tune_config=tune.TuneConfig(
|
|
search_alg=optuna_search,
|
|
),
|
|
)
|
|
tuner.fit()
|
|
|
|
# Equivalent Optuna define-by-run function approach:
|
|
|
|
def define_search_space(trial: optuna.Trial):
|
|
trial.suggest_float("a", 6, 8)
|
|
trial.suggest_float("b", 1e-4, 1e-2, log=True)
|
|
# training logic goes into trainable, this is just
|
|
# for search space definition
|
|
|
|
optuna_search = OptunaSearch(
|
|
define_search_space,
|
|
metric="loss",
|
|
mode="min")
|
|
|
|
tuner = tune.Tuner(
|
|
trainable,
|
|
tune_config=tune.TuneConfig(
|
|
search_alg=optuna_search,
|
|
),
|
|
)
|
|
tuner.fit()
|
|
|
|
Multi-objective optimization is supported:
|
|
|
|
.. code-block:: python
|
|
|
|
from ray.tune.search.optuna import OptunaSearch
|
|
import optuna
|
|
|
|
space = {
|
|
"a": optuna.distributions.FloatDistribution(6, 8),
|
|
"b": optuna.distributions.FloatDistribution(1e-4, 1e-2, log=True),
|
|
}
|
|
|
|
# Note you have to specify metric and mode here instead of
|
|
# in tune.TuneConfig
|
|
optuna_search = OptunaSearch(
|
|
space,
|
|
metric=["loss1", "loss2"],
|
|
mode=["min", "max"])
|
|
|
|
# Do not specify metric and mode here!
|
|
tuner = tune.Tuner(
|
|
trainable,
|
|
tune_config=tune.TuneConfig(
|
|
search_alg=optuna_search,
|
|
),
|
|
)
|
|
tuner.fit()
|
|
|
|
You can pass configs that will be evaluated first using
|
|
``points_to_evaluate``:
|
|
|
|
.. code-block:: python
|
|
|
|
from ray.tune.search.optuna import OptunaSearch
|
|
import optuna
|
|
|
|
space = {
|
|
"a": optuna.distributions.FloatDistribution(6, 8),
|
|
"b": optuna.distributions.FloatDistribution(1e-4, 1e-2, log=True),
|
|
}
|
|
|
|
optuna_search = OptunaSearch(
|
|
space,
|
|
points_to_evaluate=[{"a": 6.5, "b": 5e-4}, {"a": 7.5, "b": 1e-3}]
|
|
metric="loss",
|
|
mode="min")
|
|
|
|
tuner = tune.Tuner(
|
|
trainable,
|
|
tune_config=tune.TuneConfig(
|
|
search_alg=optuna_search,
|
|
),
|
|
)
|
|
tuner.fit()
|
|
|
|
Avoid re-running evaluated trials by passing the rewards together with
|
|
`points_to_evaluate`:
|
|
|
|
.. code-block:: python
|
|
|
|
from ray.tune.search.optuna import OptunaSearch
|
|
import optuna
|
|
|
|
space = {
|
|
"a": optuna.distributions.FloatDistribution(6, 8),
|
|
"b": optuna.distributions.FloatDistribution(1e-4, 1e-2, log=True),
|
|
}
|
|
|
|
optuna_search = OptunaSearch(
|
|
space,
|
|
points_to_evaluate=[{"a": 6.5, "b": 5e-4}, {"a": 7.5, "b": 1e-3}]
|
|
evaluated_rewards=[0.89, 0.42]
|
|
metric="loss",
|
|
mode="min")
|
|
|
|
tuner = tune.Tuner(
|
|
trainable,
|
|
tune_config=tune.TuneConfig(
|
|
search_alg=optuna_search,
|
|
),
|
|
)
|
|
tuner.fit()
|
|
|
|
.. versionadded:: 0.8.8
|
|
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
space: Optional[
|
|
Union[
|
|
Dict[str, "OptunaDistribution"],
|
|
List[Tuple],
|
|
Callable[["OptunaTrial"], Optional[Dict[str, Any]]],
|
|
]
|
|
] = None,
|
|
metric: Optional[Union[str, List[str]]] = None,
|
|
mode: Optional[Union[str, List[str]]] = None,
|
|
points_to_evaluate: Optional[List[Dict]] = None,
|
|
sampler: Optional["BaseSampler"] = None,
|
|
study_name: Optional[str] = None,
|
|
storage: Optional["BaseStorage"] = None,
|
|
seed: Optional[int] = None,
|
|
evaluated_rewards: Optional[List] = None,
|
|
):
|
|
assert ot is not None, "Optuna must be installed! Run `pip install optuna`."
|
|
if version.parse(ot.__version__) < version.parse("3.0.0"):
|
|
raise ImportError(
|
|
"`OptunaSearch` requires the `optuna` version to be >= 3.0.0. "
|
|
'Upgrade with: `pip install -U "optuna>=3.0"`'
|
|
)
|
|
super(OptunaSearch, self).__init__(metric=metric, mode=mode)
|
|
|
|
if isinstance(space, dict) and space:
|
|
resolved_vars, domain_vars, grid_vars = parse_spec_vars(space)
|
|
if domain_vars or grid_vars:
|
|
logger.warning(
|
|
UNRESOLVED_SEARCH_SPACE.format(par="space", cls=type(self).__name__)
|
|
)
|
|
space = self.convert_search_space(space)
|
|
else:
|
|
# Flatten to support nested dicts
|
|
space = flatten_dict(space, "/")
|
|
|
|
self._space = space
|
|
|
|
self._points_to_evaluate = points_to_evaluate or []
|
|
self._evaluated_rewards = evaluated_rewards
|
|
if study_name:
|
|
self._study_name = study_name
|
|
else:
|
|
self._study_name = "optuna" # Fixed study name for in-memory storage
|
|
|
|
if sampler and seed:
|
|
logger.warning(
|
|
"You passed an initialized sampler to `OptunaSearch`. The "
|
|
"`seed` parameter has to be passed to the sampler directly "
|
|
"and will be ignored."
|
|
)
|
|
elif sampler:
|
|
assert isinstance(sampler, BaseSampler), (
|
|
"You can only pass an instance of "
|
|
"`optuna.samplers.BaseSampler` "
|
|
"as a sampler to `OptunaSearcher`."
|
|
)
|
|
|
|
self._sampler = sampler
|
|
self._seed = seed
|
|
|
|
if storage:
|
|
assert isinstance(storage, BaseStorage), (
|
|
"The `storage` parameter in `OptunaSearcher` must be an instance "
|
|
"of `optuna.storages.BaseStorage`."
|
|
)
|
|
# If storage is not provided, just set self._storage to None
|
|
# so that the default in-memory storage is used.
|
|
self._storage = storage
|
|
|
|
self._completed_trials = set()
|
|
|
|
self._ot_trials = {}
|
|
self._ot_study = None
|
|
if self._space:
|
|
self._setup_study(mode)
|
|
|
|
def _setup_study(self, mode: Union[str, list]):
|
|
if self._metric is None and self._mode:
|
|
if isinstance(self._mode, list):
|
|
raise ValueError(
|
|
"If ``mode`` is a list (multi-objective optimization "
|
|
"case), ``metric`` must be defined."
|
|
)
|
|
# If only a mode was passed, use anonymous metric
|
|
self._metric = DEFAULT_METRIC
|
|
|
|
pruner = ot.pruners.NopPruner()
|
|
|
|
if self._sampler:
|
|
sampler = self._sampler
|
|
else:
|
|
# TPESampler handles both single- and multi-objective optimization.
|
|
sampler = ot.samplers.TPESampler(seed=self._seed)
|
|
|
|
if isinstance(mode, list):
|
|
study_direction_args = dict(
|
|
directions=["minimize" if m == "min" else "maximize" for m in mode],
|
|
)
|
|
else:
|
|
study_direction_args = dict(
|
|
direction="minimize" if mode == "min" else "maximize",
|
|
)
|
|
|
|
self._ot_study = ot.study.create_study(
|
|
storage=self._storage,
|
|
sampler=sampler,
|
|
pruner=pruner,
|
|
study_name=self._study_name,
|
|
load_if_exists=True,
|
|
**study_direction_args,
|
|
)
|
|
|
|
if self._points_to_evaluate:
|
|
validate_warmstart(
|
|
self._space,
|
|
self._points_to_evaluate,
|
|
self._evaluated_rewards,
|
|
validate_point_name_lengths=not callable(self._space),
|
|
)
|
|
if self._evaluated_rewards:
|
|
for point, reward in zip(
|
|
self._points_to_evaluate, self._evaluated_rewards
|
|
):
|
|
self.add_evaluated_point(point, reward)
|
|
else:
|
|
for point in self._points_to_evaluate:
|
|
self._ot_study.enqueue_trial(point)
|
|
|
|
def set_search_properties(
|
|
self, metric: Optional[str], mode: Optional[str], config: Dict, **spec
|
|
) -> bool:
|
|
if self._space:
|
|
return False
|
|
space = self.convert_search_space(config)
|
|
self._space = space
|
|
if metric:
|
|
self._metric = metric
|
|
if mode:
|
|
self._mode = mode
|
|
|
|
self._setup_study(self._mode)
|
|
return True
|
|
|
|
def _suggest_from_define_by_run_func(
|
|
self,
|
|
func: Callable[["OptunaTrial"], Optional[Dict[str, Any]]],
|
|
ot_trial: "OptunaTrial",
|
|
) -> Dict:
|
|
captor = _OptunaTrialSuggestCaptor(ot_trial)
|
|
time_start = time.time()
|
|
ret = func(captor)
|
|
time_taken = time.time() - time_start
|
|
if time_taken > DEFINE_BY_RUN_WARN_THRESHOLD_S:
|
|
warnings.warn(
|
|
"Define-by-run function passed in the `space` argument "
|
|
f"took {time_taken} seconds to "
|
|
"run. Ensure that actual computation, training takes "
|
|
"place inside Tune's train functions or Trainables "
|
|
"passed to `tune.Tuner()`."
|
|
)
|
|
if ret is not None:
|
|
if not isinstance(ret, dict):
|
|
raise TypeError(
|
|
"The return value of the define-by-run function "
|
|
"passed in the `space` argument should be "
|
|
"either None or a `dict` with `str` keys. "
|
|
f"Got {type(ret)}."
|
|
)
|
|
if not all(isinstance(k, str) for k in ret.keys()):
|
|
raise TypeError(
|
|
"At least one of the keys in the dict returned by the "
|
|
"define-by-run function passed in the `space` argument "
|
|
"was not a `str`."
|
|
)
|
|
return {**captor.captured_values, **ret} if ret else captor.captured_values
|
|
|
|
def suggest(self, trial_id: str) -> Optional[Dict]:
|
|
if not self._space:
|
|
raise RuntimeError(
|
|
UNDEFINED_SEARCH_SPACE.format(
|
|
cls=self.__class__.__name__, space="space"
|
|
)
|
|
)
|
|
if not self._metric or not self._mode:
|
|
raise RuntimeError(
|
|
UNDEFINED_METRIC_MODE.format(
|
|
cls=self.__class__.__name__, metric=self._metric, mode=self._mode
|
|
)
|
|
)
|
|
if callable(self._space):
|
|
# Define-by-run case
|
|
if trial_id not in self._ot_trials:
|
|
self._ot_trials[trial_id] = self._ot_study.ask()
|
|
|
|
ot_trial = self._ot_trials[trial_id]
|
|
|
|
params = self._suggest_from_define_by_run_func(self._space, ot_trial)
|
|
else:
|
|
# Use Optuna ask interface (since version 2.6.0)
|
|
if trial_id not in self._ot_trials:
|
|
self._ot_trials[trial_id] = self._ot_study.ask(
|
|
fixed_distributions=self._space
|
|
)
|
|
ot_trial = self._ot_trials[trial_id]
|
|
params = ot_trial.params
|
|
|
|
return unflatten_dict(params)
|
|
|
|
def on_trial_result(self, trial_id: str, result: Dict):
|
|
if isinstance(self.metric, list):
|
|
# Optuna doesn't support incremental results
|
|
# for multi-objective optimization
|
|
return
|
|
if trial_id in self._completed_trials:
|
|
logger.warning(
|
|
f"Received additional result for trial {trial_id}, but "
|
|
f"it already finished. Result: {result}"
|
|
)
|
|
return
|
|
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: str, result: Optional[Dict] = None, error: bool = False
|
|
):
|
|
if trial_id in self._completed_trials:
|
|
logger.warning(
|
|
f"Received additional completion for trial {trial_id}, but "
|
|
f"it already finished. Result: {result}"
|
|
)
|
|
return
|
|
|
|
ot_trial = self._ot_trials[trial_id]
|
|
|
|
if result:
|
|
if isinstance(self.metric, list):
|
|
val = [result.get(metric, None) for metric in self.metric]
|
|
else:
|
|
val = result.get(self.metric, None)
|
|
else:
|
|
val = None
|
|
ot_trial_state = OptunaTrialState.COMPLETE
|
|
if val is None:
|
|
if error:
|
|
ot_trial_state = OptunaTrialState.FAIL
|
|
else:
|
|
ot_trial_state = OptunaTrialState.PRUNED
|
|
try:
|
|
self._ot_study.tell(ot_trial, val, state=ot_trial_state)
|
|
except Exception as exc:
|
|
logger.warning(exc) # E.g. if NaN was reported
|
|
|
|
self._completed_trials.add(trial_id)
|
|
|
|
def add_evaluated_point(
|
|
self,
|
|
parameters: Dict,
|
|
value: float,
|
|
error: bool = False,
|
|
pruned: bool = False,
|
|
intermediate_values: Optional[List[float]] = None,
|
|
):
|
|
if not self._space:
|
|
raise RuntimeError(
|
|
UNDEFINED_SEARCH_SPACE.format(
|
|
cls=self.__class__.__name__, space="space"
|
|
)
|
|
)
|
|
if not self._metric or not self._mode:
|
|
raise RuntimeError(
|
|
UNDEFINED_METRIC_MODE.format(
|
|
cls=self.__class__.__name__, metric=self._metric, mode=self._mode
|
|
)
|
|
)
|
|
if callable(self._space):
|
|
raise TypeError(
|
|
"Define-by-run function passed in `space` argument is not "
|
|
"yet supported when using `evaluated_rewards`. Please provide "
|
|
"an `OptunaDistribution` dict or pass a Ray Tune "
|
|
"search space to `tune.Tuner()`."
|
|
)
|
|
|
|
ot_trial_state = OptunaTrialState.COMPLETE
|
|
if error:
|
|
ot_trial_state = OptunaTrialState.FAIL
|
|
elif pruned:
|
|
ot_trial_state = OptunaTrialState.PRUNED
|
|
|
|
if intermediate_values:
|
|
intermediate_values_dict = dict(enumerate(intermediate_values))
|
|
else:
|
|
intermediate_values_dict = None
|
|
|
|
# If the trial state is FAILED, the value must be `None` in Optuna==4.1.0
|
|
# Reference: https://github.com/optuna/optuna/pull/5211
|
|
# This is a temporary fix for the issue that Optuna enforces the value
|
|
# to be `None` if the trial state is FAILED.
|
|
# TODO (hpguo): A better solution may requires us to update the base class
|
|
# to allow the `value` arg in `add_evaluated_point` being `Optional[float]`.
|
|
if ot_trial_state == OptunaTrialState.FAIL:
|
|
value = None
|
|
|
|
trial = ot.trial.create_trial(
|
|
state=ot_trial_state,
|
|
value=value,
|
|
params=parameters,
|
|
distributions=self._space,
|
|
intermediate_values=intermediate_values_dict,
|
|
)
|
|
|
|
self._ot_study.add_trial(trial)
|
|
|
|
def save(self, checkpoint_path: str):
|
|
save_object = self.__dict__.copy()
|
|
with open(checkpoint_path, "wb") as outputFile:
|
|
pickle.dump(save_object, outputFile)
|
|
|
|
def restore(self, checkpoint_path: str):
|
|
with open(checkpoint_path, "rb") as inputFile:
|
|
save_object = pickle.load(inputFile)
|
|
if isinstance(save_object, dict):
|
|
self.__dict__.update(save_object)
|
|
else:
|
|
# Backwards compatibility
|
|
(
|
|
self._sampler,
|
|
self._ot_trials,
|
|
self._ot_study,
|
|
self._points_to_evaluate,
|
|
self._evaluated_rewards,
|
|
) = save_object
|
|
|
|
@staticmethod
|
|
def convert_search_space(spec: Dict) -> Dict[str, Any]:
|
|
resolved_vars, domain_vars, grid_vars = parse_spec_vars(spec)
|
|
|
|
if not domain_vars and not grid_vars:
|
|
return {}
|
|
|
|
if grid_vars:
|
|
raise ValueError(
|
|
"Grid search parameters cannot be automatically converted "
|
|
"to an Optuna search space."
|
|
)
|
|
|
|
# Flatten and resolve again after checking for grid search.
|
|
spec = flatten_dict(spec, prevent_delimiter=True)
|
|
resolved_vars, domain_vars, grid_vars = parse_spec_vars(spec)
|
|
|
|
def resolve_value(domain: Domain) -> ot.distributions.BaseDistribution:
|
|
quantize = None
|
|
|
|
sampler = domain.get_sampler()
|
|
if isinstance(sampler, Quantized):
|
|
quantize = sampler.q
|
|
sampler = sampler.sampler
|
|
if isinstance(sampler, LogUniform):
|
|
logger.warning(
|
|
"Optuna does not handle quantization in loguniform "
|
|
"sampling. The parameter will be passed but it will "
|
|
"probably be ignored."
|
|
)
|
|
|
|
if isinstance(domain, Float):
|
|
if isinstance(sampler, LogUniform):
|
|
if quantize:
|
|
logger.warning(
|
|
"Optuna does not support both quantization and "
|
|
"sampling from LogUniform. Dropped quantization."
|
|
)
|
|
return ot.distributions.FloatDistribution(
|
|
domain.lower, domain.upper, log=True
|
|
)
|
|
|
|
elif isinstance(sampler, Uniform):
|
|
if quantize:
|
|
return ot.distributions.FloatDistribution(
|
|
domain.lower, domain.upper, step=quantize
|
|
)
|
|
return ot.distributions.FloatDistribution(
|
|
domain.lower, domain.upper
|
|
)
|
|
|
|
elif isinstance(domain, Integer):
|
|
if isinstance(sampler, LogUniform):
|
|
return ot.distributions.IntDistribution(
|
|
domain.lower, domain.upper - 1, step=quantize or 1, log=True
|
|
)
|
|
elif isinstance(sampler, Uniform):
|
|
# Upper bound should be inclusive for quantization and
|
|
# exclusive otherwise
|
|
return ot.distributions.IntDistribution(
|
|
domain.lower,
|
|
domain.upper - int(bool(not quantize)),
|
|
step=quantize or 1,
|
|
)
|
|
elif isinstance(domain, Categorical):
|
|
if isinstance(sampler, Uniform):
|
|
return ot.distributions.CategoricalDistribution(domain.categories)
|
|
|
|
raise ValueError(
|
|
"Optuna search does not support parameters of type "
|
|
"`{}` with samplers of type `{}`".format(
|
|
type(domain).__name__, type(domain.sampler).__name__
|
|
)
|
|
)
|
|
|
|
# Parameter name is e.g. "a/b/c" for nested dicts
|
|
values = {"/".join(path): resolve_value(domain) for path, domain in domain_vars}
|
|
|
|
return values
|