Out of the box, Ax offers many options for candidate generation, most of which utilize Bayesian optimization algorithms built using BoTorch. For users that want to leverage Ax for experiment orchestration (via AxClient or Scheduler) and other features (e.g., early stopping), while relying on other methods for candidate generation, we introduced ExternalGenerationNode.
A GenerationNode is a building block of a GenerationStrategy. They can be combined together utilize different methods for generating candidates at different stages of an experiment. ExternalGenerationNode exposes a lightweight interface to allow the users to easily integrate their methods into Ax, and use them as standalone or with other GenerationNodes in a GenerationStrategy.
In this tutorial, we will implement a simple generation node using RandomForestRegressor from sklearn, and combine it with Sobol (for initialization) to optimize the Hartmann6 problem.
NOTE: This is for illustration purposes only. We do not recommend using this strategy as it typically does not perform well compared to Ax's default algorithms due to it's overly greedy behavior.
import time
from typing import Any, Dict, List, Optional, Tuple
import numpy as np
from ax.core.base_trial import TrialStatus
from ax.core.data import Data
from ax.core.experiment import Experiment
from ax.core.parameter import RangeParameter
from ax.core.types import TParameterization
from ax.modelbridge.external_generation_node import ExternalGenerationNode
from ax.modelbridge.generation_node import GenerationNode
from ax.modelbridge.generation_strategy import GenerationStrategy
from ax.modelbridge.model_spec import ModelSpec
from ax.modelbridge.registry import Models
from ax.modelbridge.transition_criterion import MaxTrials
from ax.plot.trace import plot_objective_value_vs_trial_index
from ax.service.ax_client import AxClient, ObjectiveProperties
from ax.service.utils.report_utils import exp_to_df
from ax.utils.common.typeutils import checked_cast
from ax.utils.measurement.synthetic_functions import hartmann6
from sklearn.ensemble import RandomForestRegressor
class RandomForestGenerationNode(ExternalGenerationNode):
"""A generation node that uses the RandomForestRegressor
from sklearn to predict candidate performance and picks the
next point as the random sample that has the best prediction.
To leverage external methods for candidate generation, the user must
create a subclass that implements ``update_generator_state`` and
``get_next_candidate`` methods. This can then be provided
as a node into a ``GenerationStrategy``, either as standalone or as
part of a larger generation strategy with other generation nodes,
e.g., with a Sobol node for initialization.
"""
def __init__(self, num_samples: int, regressor_options: Dict[str, Any]) -> None:
"""Initialize the generation node.
Args:
regressor_options: Options to pass to the random forest regressor.
num_samples: Number of random samples from the search space
used during candidate generation. The sample with the best
prediction is recommended as the next candidate.
"""
t_init_start = time.monotonic()
super().__init__(node_name="RandomForest")
self.num_samples: int = num_samples
self.regressor: RandomForestRegressor = RandomForestRegressor(
**regressor_options
)
# We will set these later when updating the state.
# Alternatively, we could have required experiment as an input
# and extracted them here.
self.parameters: Optional[List[RangeParameter]] = None
self.minimize: Optional[bool] = None
# Recording time spent in initializing the generator. This is
# used to compute the time spent in candidate generation.
self.fit_time_since_gen: float = time.monotonic() - t_init_start
def update_generator_state(self, experiment: Experiment, data: Data) -> None:
"""A method used to update the state of the generator. This includes any
models, predictors or any other custom state used by the generation node.
This method will be called with the up-to-date experiment and data before
``get_next_candidate`` is called to generate the next trial(s). Note
that ``get_next_candidate`` may be called multiple times (to generate
multiple candidates) after a call to ``update_generator_state``.
For this example, we will train the regressor using the latest data from
the experiment.
Args:
experiment: The ``Experiment`` object representing the current state of the
experiment. The key properties includes ``trials``, ``search_space``,
and ``optimization_config``. The data is provided as a separate arg.
data: The data / metrics collected on the experiment so far.
"""
search_space = experiment.search_space
parameter_names = list(search_space.parameters.keys())
metric_names = list(experiment.optimization_config.metrics.keys())
if any(
not isinstance(p, RangeParameter) for p in search_space.parameters.values()
):
raise NotImplementedError(
"This example only supports RangeParameters in the search space."
)
if search_space.parameter_constraints:
raise NotImplementedError(
"This example does not support parameter constraints."
)
if len(metric_names) != 1:
raise NotImplementedError(
"This example only supports single-objective optimization."
)
# Get the data for the completed trials.
num_completed_trials = len(experiment.trials_by_status[TrialStatus.COMPLETED])
x = np.zeros([num_completed_trials, len(parameter_names)])
y = np.zeros([num_completed_trials, 1])
for t_idx, trial in experiment.trials.items():
if trial.status == "COMPLETED":
trial_parameters = trial.arm.parameters
x[t_idx, :] = np.array([trial_parameters[p] for p in parameter_names])
trial_df = data.df[data.df["trial_index"] == t_idx]
y[t_idx, 0] = trial_df[trial_df["metric_name"] == metric_names[0]][
"mean"
].item()
# Train the regressor.
self.regressor.fit(x, y)
# Update the attributes not set in __init__.
self.parameters = search_space.parameters
self.minimize = experiment.optimization_config.objective.minimize
def get_next_candidate(
self, pending_parameters: List[TParameterization]
) -> TParameterization:
"""Get the parameters for the next candidate configuration to evaluate.
We will draw ``self.num_samples`` random samples from the search space
and predict the objective value for each sample. We will then return
the sample with the best predicted value.
Args:
pending_parameters: A list of parameters of the candidates pending
evaluation. This is often used to avoid generating duplicate candidates.
We ignore this here for simplicity.
Returns:
A dictionary mapping parameter names to parameter values for the next
candidate suggested by the method.
"""
bounds = np.array([[p.lower, p.upper] for p in self.parameters.values()])
unit_samples = np.random.random_sample([self.num_samples, len(bounds)])
samples = bounds[:, 0] + (bounds[:, 1] - bounds[:, 0]) * unit_samples
# Predict the objective value for each sample.
y_pred = self.regressor.predict(samples)
# Find the best sample.
best_idx = np.argmin(y_pred) if self.minimize else np.argmax(y_pred)
best_sample = samples[best_idx, :]
# Convert the sample to a parameterization.
candidate = {
p_name: best_sample[i].item()
for i, p_name in enumerate(self.parameters.keys())
}
return candidate
We will use Sobol for the first 5 trials and defer to random forest for the rest.
generation_strategy = GenerationStrategy(
name="Sobol+RandomForest",
nodes=[
GenerationNode(
node_name="Sobol",
model_specs=[ModelSpec(Models.SOBOL)],
transition_criteria=[
MaxTrials(
# This specifies the maximum number of trials to generate from this node,
# and the next node in the strategy.
threshold=5,
block_transition_if_unmet=True,
transition_to="RandomForest"
)
],
),
RandomForestGenerationNode(num_samples=128, regressor_options={}),
],
)
ax_client = AxClient(generation_strategy=generation_strategy)
ax_client.create_experiment(
name="hartmann_test_experiment",
parameters=[
{
"name": f"x{i}",
"type": "range",
"bounds": [0.0, 1.0],
"value_type": "float", # Optional, defaults to inference from type of "bounds".
}
for i in range(1, 7)
],
objectives={"hartmann6": ObjectiveProperties(minimize=True)},
)
def evaluate(parameterization: TParameterization) -> Dict[str, Tuple[float, float]]:
x = np.array([parameterization.get(f"x{i+1}") for i in range(6)])
return {"hartmann6": (checked_cast(float, hartmann6(x)), 0.0)}
[INFO 07-23 21:43:59] ax.service.ax_client: Starting optimization with verbose logging. To disable logging, set the `verbose_logging` argument to `False`. Note that float values in the logs are rounded to 6 decimal points.
[INFO 07-23 21:43:59] ax.service.utils.instantiation: Created search space: SearchSpace(parameters=[RangeParameter(name='x1', parameter_type=FLOAT, range=[0.0, 1.0]), RangeParameter(name='x2', parameter_type=FLOAT, range=[0.0, 1.0]), RangeParameter(name='x3', parameter_type=FLOAT, range=[0.0, 1.0]), RangeParameter(name='x4', parameter_type=FLOAT, range=[0.0, 1.0]), RangeParameter(name='x5', parameter_type=FLOAT, range=[0.0, 1.0]), RangeParameter(name='x6', parameter_type=FLOAT, range=[0.0, 1.0])], parameter_constraints=[]).
for i in range(15):
parameterization, trial_index = ax_client.get_next_trial()
ax_client.complete_trial(
trial_index=trial_index, raw_data=evaluate(parameterization)
)
/tmp/tmp.DL1QmpHQMI/Ax-main/ax/modelbridge/cross_validation.py:462: UserWarning: Encountered exception in computing model fit quality: RandomModelBridge does not support prediction.
warn("Encountered exception in computing model fit quality: " + str(e))
[INFO 07-23 21:43:59] ax.service.ax_client: Generated new trial 0 with parameters {'x1': 0.235943, 'x2': 0.805732, 'x3': 0.851269, 'x4': 0.025475, 'x5': 0.219072, 'x6': 0.733228} using model Sobol.
[INFO 07-23 21:43:59] ax.service.ax_client: Completed trial 0 with data: {'hartmann6': (-0.293255, 0.0)}.
/tmp/tmp.DL1QmpHQMI/Ax-main/ax/modelbridge/cross_validation.py:462: UserWarning: Encountered exception in computing model fit quality: RandomModelBridge does not support prediction.
warn("Encountered exception in computing model fit quality: " + str(e))
[INFO 07-23 21:43:59] ax.service.ax_client: Generated new trial 1 with parameters {'x1': 0.070086, 'x2': 0.988092, 'x3': 0.86912, 'x4': 0.609664, 'x5': 0.076967, 'x6': 0.834669} using model Sobol.
[INFO 07-23 21:43:59] ax.service.ax_client: Completed trial 1 with data: {'hartmann6': (-0.048005, 0.0)}.
/tmp/tmp.DL1QmpHQMI/Ax-main/ax/modelbridge/cross_validation.py:462: UserWarning: Encountered exception in computing model fit quality: RandomModelBridge does not support prediction.
warn("Encountered exception in computing model fit quality: " + str(e))
[INFO 07-23 21:43:59] ax.service.ax_client: Generated new trial 2 with parameters {'x1': 0.135788, 'x2': 0.209062, 'x3': 0.116069, 'x4': 0.590754, 'x5': 0.331103, 'x6': 0.194263} using model Sobol.
[INFO 07-23 21:43:59] ax.service.ax_client: Completed trial 2 with data: {'hartmann6': (-0.194781, 0.0)}.
/tmp/tmp.DL1QmpHQMI/Ax-main/ax/modelbridge/cross_validation.py:462: UserWarning: Encountered exception in computing model fit quality: RandomModelBridge does not support prediction.
warn("Encountered exception in computing model fit quality: " + str(e))
[INFO 07-23 21:43:59] ax.service.ax_client: Generated new trial 3 with parameters {'x1': 0.797577, 'x2': 0.627523, 'x3': 0.322616, 'x4': 0.869181, 'x5': 0.80381, 'x6': 0.54154} using model Sobol.
[INFO 07-23 21:43:59] ax.service.ax_client: Completed trial 3 with data: {'hartmann6': (-0.00198, 0.0)}.
/tmp/tmp.DL1QmpHQMI/Ax-main/ax/modelbridge/cross_validation.py:462: UserWarning: Encountered exception in computing model fit quality: RandomModelBridge does not support prediction.
warn("Encountered exception in computing model fit quality: " + str(e))
[INFO 07-23 21:43:59] ax.service.ax_client: Generated new trial 4 with parameters {'x1': 0.276039, 'x2': 0.370268, 'x3': 0.655317, 'x4': 0.068835, 'x5': 0.77193, 'x6': 0.346295} using model Sobol.
[INFO 07-23 21:43:59] ax.service.ax_client: Completed trial 4 with data: {'hartmann6': (-0.394922, 0.0)}.
/opt/hostedtoolcache/Python/3.10.14/x64/lib/python3.10/site-packages/sklearn/base.py:1473: DataConversionWarning: A column-vector y was passed when a 1d array was expected. Please change the shape of y to (n_samples,), for example using ravel().
return fit_method(estimator, *args, **kwargs)
[INFO 07-23 21:43:59] ax.service.ax_client: Generated new trial 5 with parameters {'x1': 0.724667, 'x2': 0.171867, 'x3': 0.389522, 'x4': 0.268447, 'x5': 0.176725, 'x6': 0.214398} using model RandomForest.
[INFO 07-23 21:43:59] ax.service.ax_client: Completed trial 5 with data: {'hartmann6': (-0.220702, 0.0)}.
/opt/hostedtoolcache/Python/3.10.14/x64/lib/python3.10/site-packages/sklearn/base.py:1473: DataConversionWarning: A column-vector y was passed when a 1d array was expected. Please change the shape of y to (n_samples,), for example using ravel(). return fit_method(estimator, *args, **kwargs)
[INFO 07-23 21:43:59] ax.service.ax_client: Generated new trial 6 with parameters {'x1': 0.590404, 'x2': 0.434876, 'x3': 0.581286, 'x4': 0.06917, 'x5': 0.266751, 'x6': 0.790401} using model RandomForest.
[INFO 07-23 21:43:59] ax.service.ax_client: Completed trial 6 with data: {'hartmann6': (-0.963571, 0.0)}.
/opt/hostedtoolcache/Python/3.10.14/x64/lib/python3.10/site-packages/sklearn/base.py:1473: DataConversionWarning: A column-vector y was passed when a 1d array was expected. Please change the shape of y to (n_samples,), for example using ravel(). return fit_method(estimator, *args, **kwargs)
[INFO 07-23 21:43:59] ax.service.ax_client: Generated new trial 7 with parameters {'x1': 0.159938, 'x2': 0.59064, 'x3': 0.059199, 'x4': 0.142424, 'x5': 0.740353, 'x6': 0.668626} using model RandomForest.
[INFO 07-23 21:43:59] ax.service.ax_client: Completed trial 7 with data: {'hartmann6': (-0.048791, 0.0)}.
/opt/hostedtoolcache/Python/3.10.14/x64/lib/python3.10/site-packages/sklearn/base.py:1473: DataConversionWarning: A column-vector y was passed when a 1d array was expected. Please change the shape of y to (n_samples,), for example using ravel(). return fit_method(estimator, *args, **kwargs)
[INFO 07-23 21:43:59] ax.service.ax_client: Generated new trial 8 with parameters {'x1': 0.349523, 'x2': 0.827577, 'x3': 0.078603, 'x4': 0.731062, 'x5': 0.383811, 'x6': 0.956447} using model RandomForest.
[INFO 07-23 21:43:59] ax.service.ax_client: Completed trial 8 with data: {'hartmann6': (-0.032078, 0.0)}.
/opt/hostedtoolcache/Python/3.10.14/x64/lib/python3.10/site-packages/sklearn/base.py:1473: DataConversionWarning: A column-vector y was passed when a 1d array was expected. Please change the shape of y to (n_samples,), for example using ravel(). return fit_method(estimator, *args, **kwargs)
[INFO 07-23 21:43:59] ax.service.ax_client: Generated new trial 9 with parameters {'x1': 0.386376, 'x2': 0.847778, 'x3': 0.594479, 'x4': 0.472761, 'x5': 0.161133, 'x6': 0.253349} using model RandomForest.
[INFO 07-23 21:43:59] ax.service.ax_client: Completed trial 9 with data: {'hartmann6': (-1.549365, 0.0)}.
/opt/hostedtoolcache/Python/3.10.14/x64/lib/python3.10/site-packages/sklearn/base.py:1473: DataConversionWarning: A column-vector y was passed when a 1d array was expected. Please change the shape of y to (n_samples,), for example using ravel(). return fit_method(estimator, *args, **kwargs)
[INFO 07-23 21:43:59] ax.service.ax_client: Generated new trial 10 with parameters {'x1': 0.337886, 'x2': 0.127233, 'x3': 0.946171, 'x4': 0.161489, 'x5': 0.152265, 'x6': 0.984759} using model RandomForest.
[INFO 07-23 21:43:59] ax.service.ax_client: Completed trial 10 with data: {'hartmann6': (-0.818642, 0.0)}.
/opt/hostedtoolcache/Python/3.10.14/x64/lib/python3.10/site-packages/sklearn/base.py:1473: DataConversionWarning: A column-vector y was passed when a 1d array was expected. Please change the shape of y to (n_samples,), for example using ravel(). return fit_method(estimator, *args, **kwargs)
[INFO 07-23 21:43:59] ax.service.ax_client: Generated new trial 11 with parameters {'x1': 0.597673, 'x2': 0.537541, 'x3': 0.132777, 'x4': 0.331666, 'x5': 0.252365, 'x6': 0.765986} using model RandomForest.
[INFO 07-23 21:43:59] ax.service.ax_client: Completed trial 11 with data: {'hartmann6': (-0.939178, 0.0)}.
/opt/hostedtoolcache/Python/3.10.14/x64/lib/python3.10/site-packages/sklearn/base.py:1473: DataConversionWarning: A column-vector y was passed when a 1d array was expected. Please change the shape of y to (n_samples,), for example using ravel(). return fit_method(estimator, *args, **kwargs)
[INFO 07-23 21:43:59] ax.service.ax_client: Generated new trial 12 with parameters {'x1': 0.150772, 'x2': 0.953183, 'x3': 0.963038, 'x4': 0.251849, 'x5': 0.854956, 'x6': 0.759404} using model RandomForest.
[INFO 07-23 21:43:59] ax.service.ax_client: Completed trial 12 with data: {'hartmann6': (-0.007497, 0.0)}.
/opt/hostedtoolcache/Python/3.10.14/x64/lib/python3.10/site-packages/sklearn/base.py:1473: DataConversionWarning: A column-vector y was passed when a 1d array was expected. Please change the shape of y to (n_samples,), for example using ravel(). return fit_method(estimator, *args, **kwargs)
[INFO 07-23 21:44:00] ax.service.ax_client: Generated new trial 13 with parameters {'x1': 0.080015, 'x2': 0.897607, 'x3': 0.291302, 'x4': 0.73702, 'x5': 0.169201, 'x6': 0.291463} using model RandomForest.
[INFO 07-23 21:44:00] ax.service.ax_client: Completed trial 13 with data: {'hartmann6': (-0.17789, 0.0)}.
/opt/hostedtoolcache/Python/3.10.14/x64/lib/python3.10/site-packages/sklearn/base.py:1473: DataConversionWarning: A column-vector y was passed when a 1d array was expected. Please change the shape of y to (n_samples,), for example using ravel(). return fit_method(estimator, *args, **kwargs)
[INFO 07-23 21:44:00] ax.service.ax_client: Generated new trial 14 with parameters {'x1': 0.073141, 'x2': 0.482493, 'x3': 0.34948, 'x4': 0.528698, 'x5': 0.087757, 'x6': 0.710098} using model RandomForest.
[INFO 07-23 21:44:00] ax.service.ax_client: Completed trial 14 with data: {'hartmann6': (-0.516728, 0.0)}.
exp_df = exp_to_df(ax_client.experiment)
exp_df
[WARNING 07-23 21:44:00] ax.service.utils.report_utils: Column reason missing for all trials. Not appending column.
| trial_index | arm_name | trial_status | generation_method | hartmann6 | x1 | x2 | x3 | x4 | x5 | x6 | |
|---|---|---|---|---|---|---|---|---|---|---|---|
| 0 | 0 | 0_0 | COMPLETED | Sobol | -0.293255 | 0.235943 | 0.805732 | 0.851269 | 0.025475 | 0.219072 | 0.733228 |
| 1 | 1 | 1_0 | COMPLETED | Sobol | -0.048005 | 0.070086 | 0.988092 | 0.869120 | 0.609664 | 0.076967 | 0.834669 |
| 2 | 2 | 2_0 | COMPLETED | Sobol | -0.194781 | 0.135788 | 0.209062 | 0.116069 | 0.590754 | 0.331103 | 0.194263 |
| 3 | 3 | 3_0 | COMPLETED | Sobol | -0.001980 | 0.797577 | 0.627523 | 0.322616 | 0.869181 | 0.803810 | 0.541540 |
| 4 | 4 | 4_0 | COMPLETED | Sobol | -0.394922 | 0.276039 | 0.370268 | 0.655317 | 0.068835 | 0.771930 | 0.346295 |
| 5 | 5 | 5_0 | COMPLETED | RandomForest | -0.220702 | 0.724667 | 0.171867 | 0.389522 | 0.268447 | 0.176725 | 0.214398 |
| 6 | 6 | 6_0 | COMPLETED | RandomForest | -0.963571 | 0.590404 | 0.434876 | 0.581286 | 0.069170 | 0.266751 | 0.790401 |
| 7 | 7 | 7_0 | COMPLETED | RandomForest | -0.048791 | 0.159938 | 0.590640 | 0.059199 | 0.142424 | 0.740353 | 0.668626 |
| 8 | 8 | 8_0 | COMPLETED | RandomForest | -0.032078 | 0.349523 | 0.827577 | 0.078603 | 0.731062 | 0.383811 | 0.956447 |
| 9 | 9 | 9_0 | COMPLETED | RandomForest | -1.549365 | 0.386376 | 0.847778 | 0.594479 | 0.472761 | 0.161133 | 0.253349 |
| 10 | 10 | 10_0 | COMPLETED | RandomForest | -0.818642 | 0.337886 | 0.127233 | 0.946171 | 0.161489 | 0.152265 | 0.984759 |
| 11 | 11 | 11_0 | COMPLETED | RandomForest | -0.939178 | 0.597673 | 0.537541 | 0.132777 | 0.331666 | 0.252365 | 0.765986 |
| 12 | 12 | 12_0 | COMPLETED | RandomForest | -0.007497 | 0.150772 | 0.953183 | 0.963038 | 0.251849 | 0.854956 | 0.759404 |
| 13 | 13 | 13_0 | COMPLETED | RandomForest | -0.177890 | 0.080015 | 0.897607 | 0.291302 | 0.737020 | 0.169201 | 0.291463 |
| 14 | 14 | 14_0 | COMPLETED | RandomForest | -0.516728 | 0.073141 | 0.482493 | 0.349480 | 0.528698 | 0.087757 | 0.710098 |
plot_objective_value_vs_trial_index(
exp_df=exp_df,
metric_colname="hartmann6",
minimize=True,
title="Hartmann6 Objective Value vs. Trial Index",
)
/tmp/tmp.DL1QmpHQMI/Ax-main/ax/plot/trace.py:852: FutureWarning: DataFrame.fillna with 'method' is deprecated and will raise in a future version. Use obj.ffill() or obj.bfill() instead.
Total runtime of script: 8.09 seconds.