Source code for ray.tune.suggest.ax

from typing import Dict

from ax.service.ax_client import AxClient
from ray.tune.sample import Categorical, Float, Integer, LogUniform, \
    Quantized, Uniform
from ray.tune.suggest.variant_generator import parse_spec_vars
from ray.tune.utils import flatten_dict
from ray.tune.utils.util import unflatten_dict

try:
    import ax
except ImportError:
    ax = None
import logging

from ray.tune.suggest import Searcher

logger = logging.getLogger(__name__)


[docs]class AxSearch(Searcher): """Uses `Ax <https://ax.dev/>`_ to optimize hyperparameters. Ax is a platform for understanding, managing, deploying, and automating adaptive experiments. Ax provides an easy to use interface with BoTorch, a flexible, modern library for Bayesian optimization in PyTorch. More information can be found in https://ax.dev/. To use this search algorithm, you must install Ax and sqlalchemy: .. code-block:: bash $ pip install ax-platform sqlalchemy Parameters: space (list[dict]): Parameters in the experiment search space. Required elements in the dictionaries are: "name" (name of this parameter, string), "type" (type of the parameter: "range", "fixed", or "choice", string), "bounds" for range parameters (list of two values, lower bound first), "values" for choice parameters (list of values), and "value" for fixed parameters (single value). objective_name (str): Name of the metric used as objective in this experiment. This metric must be present in `raw_data` argument to `log_data`. This metric must also be present in the dict reported/returned by the Trainable. mode (str): One of {min, max}. Determines whether objective is minimizing or maximizing the metric attribute. Defaults to "max". parameter_constraints (list[str]): Parameter constraints, such as "x3 >= x4" or "x3 + x4 >= 2". outcome_constraints (list[str]): Outcome constraints of form "metric_name >= bound", like "m1 <= 3." ax_client (AxClient): Optional AxClient instance. If this is set, do not pass any values to these parameters: `space`, `objective_name`, `parameter_constraints`, `outcome_constraints`. use_early_stopped_trials: Deprecated. max_concurrent (int): Deprecated. Tune automatically converts search spaces to Ax's format: .. code-block:: python from ray import tune from ray.tune.suggest.ax import AxSearch config = { "x1": tune.uniform(0.0, 1.0), "x2": tune.uniform(0.0, 1.0) } def easy_objective(config): for i in range(100): intermediate_result = config["x1"] + config["x2"] * i tune.report(score=intermediate_result) ax_search = AxSearch(objective_name="score") tune.run( config=config, easy_objective, search_alg=ax_search) If you would like to pass the search space manually, the code would look like this: .. code-block:: python from ray import tune from ray.tune.suggest.ax import AxSearch parameters = [ {"name": "x1", "type": "range", "bounds": [0.0, 1.0]}, {"name": "x2", "type": "range", "bounds": [0.0, 1.0]}, ] def easy_objective(config): for i in range(100): intermediate_result = config["x1"] + config["x2"] * i tune.report(score=intermediate_result) ax_search = AxSearch(space=parameters, objective_name="score") tune.run(easy_objective, search_alg=ax_search) """ def __init__(self, space=None, metric="episode_reward_mean", mode="max", parameter_constraints=None, outcome_constraints=None, ax_client=None, use_early_stopped_trials=None, max_concurrent=None): assert ax is not None, "Ax must be installed!" assert mode in ["min", "max"], "`mode` must be one of ['min', 'max']" super(AxSearch, self).__init__( metric=metric, mode=mode, max_concurrent=max_concurrent, use_early_stopped_trials=use_early_stopped_trials) self._ax = ax_client self._space = space self._parameter_constraints = parameter_constraints self._outcome_constraints = outcome_constraints self.max_concurrent = max_concurrent self._objective_name = metric self._parameters = [] self._live_trial_mapping = {} if self._ax or self._space: self.setup_experiment() def setup_experiment(self): if not self._ax: self._ax = AxClient() try: exp = self._ax.experiment has_experiment = True except ValueError: has_experiment = False if not has_experiment: if not self._space: raise ValueError( "You have to create an Ax experiment by calling " "`AxClient.create_experiment()`, or you should pass an " "Ax search space as the `space` parameter to `AxSearch`, " "or pass a `config` dict to `tune.run()`.") self._ax.create_experiment( parameters=self._space, objective_name=self._metric, parameter_constraints=self._parameter_constraints, outcome_constraints=self._outcome_constraints, minimize=self._mode != "max") else: if any([ self._space, self._parameter_constraints, self._outcome_constraints ]): raise ValueError( "If you create the Ax experiment yourself, do not pass " "values for these parameters to `AxSearch`: {}.".format([ "space", "parameter_constraints", "outcome_constraints" ])) exp = self._ax.experiment self._objective_name = exp.optimization_config.objective.metric.name self._parameters = list(exp.parameters) if self._ax._enforce_sequential_optimization: logger.warning("Detected sequential enforcement. Be sure to use " "a ConcurrencyLimiter.") def set_search_properties(self, metric, mode, config): if self._ax: return False space = self.convert_search_space(config) self._space = space if metric: self._metric = metric if mode: self._mode = mode self.setup_experiment() return True def suggest(self, trial_id): if not self._ax: raise RuntimeError( "Trying to sample a configuration from {}, but no search " "space has been defined. Either pass the `{}` argument when " "instantiating the search algorithm, or pass a `config` to " "`tune.run()`.".format(self.__class__.__name__, "space")) if self.max_concurrent: if len(self._live_trial_mapping) >= self.max_concurrent: return None parameters, trial_index = self._ax.get_next_trial() self._live_trial_mapping[trial_id] = trial_index return unflatten_dict(parameters) def on_trial_complete(self, trial_id, result=None, error=False): """Notification for the completion of trial. Data of form key value dictionary of metric names and values. """ if result: self._process_result(trial_id, result) self._live_trial_mapping.pop(trial_id) def _process_result(self, trial_id, result): ax_trial_index = self._live_trial_mapping[trial_id] metric_dict = { self._objective_name: (result[self._objective_name], 0.0) } outcome_names = [ oc.metric.name for oc in self._ax.experiment.optimization_config.outcome_constraints ] metric_dict.update({on: (result[on], 0.0) for on in outcome_names}) self._ax.complete_trial( trial_index=ax_trial_index, raw_data=metric_dict) @staticmethod def convert_search_space(spec: Dict): spec = flatten_dict(spec, prevent_delimiter=True) resolved_vars, domain_vars, grid_vars = parse_spec_vars(spec) if grid_vars: raise ValueError( "Grid search parameters cannot be automatically converted " "to an Ax search space.") def resolve_value(par, domain): sampler = domain.get_sampler() if isinstance(sampler, Quantized): logger.warning("AxSearch does not support quantization. " "Dropped quantization.") sampler = sampler.sampler if isinstance(domain, Float): if isinstance(sampler, LogUniform): return { "name": par, "type": "range", "bounds": [domain.lower, domain.upper], "value_type": "float", "log_scale": True } elif isinstance(sampler, Uniform): return { "name": par, "type": "range", "bounds": [domain.lower, domain.upper], "value_type": "float", "log_scale": False } elif isinstance(domain, Integer): if isinstance(sampler, LogUniform): return { "name": par, "type": "range", "bounds": [domain.lower, domain.upper], "value_type": "int", "log_scale": True } elif isinstance(sampler, Uniform): return { "name": par, "type": "range", "bounds": [domain.lower, domain.upper], "value_type": "int", "log_scale": False } elif isinstance(domain, Categorical): if isinstance(sampler, Uniform): return { "name": par, "type": "choice", "values": domain.categories } raise ValueError("AxSearch does not support parameters of type " "`{}` with samplers of type `{}`".format( type(domain).__name__, type(domain.sampler).__name__)) # Fixed vars fixed_values = [{ "name": "/".join(path), "type": "fixed", "value": val } for path, val in resolved_vars] # Parameter name is e.g. "a/b/c" for nested dicts resolved_values = [ resolve_value("/".join(path), domain) for path, domain in domain_vars ] return fixed_values + resolved_values