ray.tune.search.ax.AxSearch#

class ray.tune.search.ax.AxSearch(space: Optional[Union[Dict, List[Dict]]] = None, metric: Optional[str] = None, mode: Optional[str] = None, points_to_evaluate: Optional[List[Dict]] = None, parameter_constraints: Optional[List] = None, outcome_constraints: Optional[List] = None, ax_client: None = None, **ax_kwargs)[source]#

Bases: ray.tune.search.searcher.Searcher

Uses Ax 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:

$ pip install ax-platform sqlalchemy
Parameters

  • space – 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).

  • metric – 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. If None but a mode was passed, the ray.tune.result.DEFAULT_METRIC will be used per default.

  • mode – One of {min, max}. Determines whether objective is minimizing or maximizing the metric attribute. Defaults to “max”.

  • 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.

  • parameter_constraints – Parameter constraints, such as “x3 >= x4” or “x3 + x4 >= 2”.

  • outcome_constraints – Outcome constraints of form “metric_name >= bound”, like “m1 <= 3.”

  • ax_client – Optional AxClient instance. If this is set, do not pass any values to these parameters: space, metric, parameter_constraints, outcome_constraints.

  • **ax_kwargs – Passed to AxClient instance. Ignored if AxClient is not None.

Tune automatically converts search spaces to Ax’s format:

from ray import train, tune
from ray.tune.search.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
        train.report({"score": intermediate_result})

ax_search = AxSearch()
tuner = tune.Tuner(
    easy_objective,
    tune_config=tune.TuneConfig(
        search_alg=ax_search,
        metric="score",
        mode="max",
    ),
    param_space=config,
)
tuner.fit()

If you would like to pass the search space manually, the code would look like this:

from ray import train, tune
from ray.tune.search.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
        train.report({"score": intermediate_result})

ax_search = AxSearch(space=parameters, metric="score", mode="max")
tuner = tune.Tuner(
    easy_objective,
    tune_config=tune.TuneConfig(
        search_alg=ax_search,
    ),
)
tuner.fit()

Methods

add_evaluated_point(parameters, value[, ...])

Pass results from a point that has been evaluated separately.

add_evaluated_trials(trials_or_analysis, metric)

Pass results from trials that have been evaluated separately.

on_trial_complete(trial_id[, result, error])

Notification for the completion of trial.

on_trial_result(trial_id, result)

Optional notification for result during training.

restore_from_dir(checkpoint_dir)

Restores the state of a searcher from a given checkpoint_dir.

save_to_dir(checkpoint_dir[, session_str])

Automatically saves the given searcher to the checkpoint_dir.

set_max_concurrency(max_concurrent)

Set max concurrent trials this searcher can run.

Attributes

CKPT_FILE_TMPL

FINISHED

metric

The training result objective value attribute.

mode

Specifies if minimizing or maximizing the metric.