Skip to main content
Version: 0.5.0

Tune a CNN on MNIST

This tutorial walks through using Ax to tune two hyperparameters (learning rate and momentum) for a PyTorch CNN on the MNIST dataset trained using SGD with momentum.

import sys
in_colab = 'google.colab' in sys.modules
if in_colab:
%pip install ax-platform
import torch

import torch.nn as nn
import torch.nn.functional as F

from ax.service.ax_client import AxClient, ObjectiveProperties
from ax.service.utils.report_utils import exp_to_df
from ax.utils.notebook.plotting import init_notebook_plotting, render
from ax.utils.tutorials.cnn_utils import evaluate, load_mnist, train
from torch._tensor import Tensor
from torch.utils.data import DataLoader
import plotly.io as pio

init_notebook_plotting()
if in_colab:
pio.renderers.default = "colab"
Out:

[INFO 02-03 18:40:11] ax.utils.notebook.plotting: Injecting Plotly library into cell. Do not overwrite or delete cell.

Out:

[INFO 02-03 18:40:11] ax.utils.notebook.plotting: Please see

(https://ax.dev/tutorials/visualizations.html#Fix-for-plots-that-are-not-rendering)

if visualizations are not rendering.

torch.manual_seed(42)
dtype = torch.float
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

1. Load MNIST data

First, we need to load the MNIST data and partition it into training, validation, and test sets.

Note: this will download the dataset if necessary.

BATCH_SIZE = 512
train_loader, valid_loader, test_loader = load_mnist(batch_size=BATCH_SIZE)
Out:

0%| | 0.00/9.91M [00:00<?, ?B/s]

Out:

12%|█▏ | 1.15M/9.91M [00:00<00:00, 11.5MB/s]

Out:

100%|██████████| 9.91M/9.91M [00:00<00:00, 54.1MB/s]

Out:

0%| | 0.00/28.9k [00:00<?, ?B/s]

Out:

100%|██████████| 28.9k/28.9k [00:00<00:00, 1.84MB/s]

Out:

0%| | 0.00/1.65M [00:00<?, ?B/s]

Out:

70%|██████▉ | 1.15M/1.65M [00:00<00:00, 11.4MB/s]

Out:

100%|██████████| 1.65M/1.65M [00:00<00:00, 16.1MB/s]

Out:

0%| | 0.00/4.54k [00:00<?, ?B/s]

Out:

100%|██████████| 4.54k/4.54k [00:00<00:00, 10.5MB/s]

2. Initialize Client

Create a client object to interface with Ax APIs. By default this runs locally without storage.

ax_client = AxClient()
Out:

[INFO 02-03 18:40:13] 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.

3. Set up experiment

An experiment consists of a search space (parameters and parameter constraints) and optimization configuration (objective name, minimization setting, and outcome constraints).

# Create an experiment with required arguments: name, parameters, and objective_name.
ax_client.create_experiment(
name="tune_cnn_on_mnist", # The name of the experiment.
parameters=[
{
"name": "lr", # The name of the parameter.
"type": "range", # The type of the parameter ("range", "choice" or "fixed").
"bounds": [1e-6, 0.4], # The bounds for range parameters.
# "values" The possible values for choice parameters .
# "value" The fixed value for fixed parameters.
"value_type": "float", # Optional, the value type ("int", "float", "bool" or "str"). Defaults to inference from type of "bounds".
"log_scale": True, # Optional, whether to use a log scale for range parameters. Defaults to False.
# "is_ordered" Optional, a flag for choice parameters.
},
{
"name": "momentum",
"type": "range",
"bounds": [0.0, 1.0],
},
],
objectives={"accuracy": ObjectiveProperties(minimize=False)}, # The objective name and minimization setting.
# parameter_constraints: Optional, a list of strings of form "p1 >= p2" or "p1 + p2 <= some_bound".
# outcome_constraints: Optional, a list of strings of form "constrained_metric <= some_bound".
)
Out:

[INFO 02-03 18:40:14] ax.service.utils.instantiation: Inferred value type of ParameterType.FLOAT for parameter momentum. If that is not the expected value type, you can explicitly specify 'value_type' ('int', 'float', 'bool' or 'str') in parameter dict.

Out:

[INFO 02-03 18:40:14] ax.service.utils.instantiation: Created search space: SearchSpace(parameters=[RangeParameter(name='lr', parameter_type=FLOAT, range=[1e-06, 0.4], log_scale=True), RangeParameter(name='momentum', parameter_type=FLOAT, range=[0.0, 1.0])], parameter_constraints=[]).

Out:

[INFO 02-03 18:40:14] ax.modelbridge.dispatch_utils: Using Models.BOTORCH_MODULAR since there is at least one ordered parameter and there are no unordered categorical parameters.

Out:

[INFO 02-03 18:40:14] ax.modelbridge.dispatch_utils: Calculating the number of remaining initialization trials based on num_initialization_trials=None max_initialization_trials=None num_tunable_parameters=2 num_trials=None use_batch_trials=False

Out:

[INFO 02-03 18:40:14] ax.modelbridge.dispatch_utils: calculated num_initialization_trials=5

Out:

[INFO 02-03 18:40:14] ax.modelbridge.dispatch_utils: num_completed_initialization_trials=0 num_remaining_initialization_trials=5

Out:

[INFO 02-03 18:40:14] ax.modelbridge.dispatch_utils: verbose, disable_progbar, and jit_compile are not yet supported when using choose_generation_strategy with ModularBoTorchModel, dropping these arguments.

Out:

[INFO 02-03 18:40:14] ax.modelbridge.dispatch_utils: Using Bayesian Optimization generation strategy: GenerationStrategy(name='Sobol+BoTorch', steps=[Sobol for 5 trials, BoTorch for subsequent trials]). Iterations after 5 will take longer to generate due to model-fitting.

4. Define how to evaluate trials

First we define a simple CNN class to classify the MNIST images

class CNN(nn.Module):

def __init__(self) -> None:
super().__init__()
self.conv1 = nn.Conv2d(1, 20, kernel_size=5, stride=1)
self.fc1 = nn.Linear(8 * 8 * 20, 64)
self.fc2 = nn.Linear(64, 10)

def forward(self, x: Tensor) -> Tensor:
x = F.relu(self.conv1(x))
x = F.max_pool2d(x, 3, 3)
x = x.view(-1, 8 * 8 * 20)
x = F.relu(self.fc1(x))
x = self.fc2(x)
return F.log_softmax(x, dim=-1)

In this tutorial, we want to optimize classification accuracy on the validation set as a function of the learning rate and momentum. The train_evaluate function takes in a parameterization (set of parameter values), computes the classification accuracy, and returns that metric.

def train_evaluate(parameterization):
"""
Train the model and then compute an evaluation metric.

In this tutorial, the CNN utils package is doing a lot of work
under the hood:
- `train` initializes the network, defines the loss function
and optimizer, performs the training loop, and returns the
trained model.
- `evaluate` computes the accuracy of the model on the
evaluation dataset and returns the metric.

For your use case, you can define training and evaluation functions
of your choosing.

"""
net = CNN()
net = train(
net=net,
train_loader=train_loader,
parameters=parameterization,
dtype=dtype,
device=device,
)

return evaluate(
net=net,
data_loader=valid_loader,
dtype=dtype,
device=device,
)

5. Run optimization loop

First we use attach_trial to attach a custom trial with manually-chosen parameters. This step is optional, but we include it here to demonstrate adding manual trials and to serve as a baseline model with decent performance.

# Attach the trial
ax_client.attach_trial(
parameters={"lr": 0.000026, "momentum": 0.58}
)

# Get the parameters and run the trial
baseline_parameters = ax_client.get_trial_parameters(trial_index=0)
ax_client.complete_trial(trial_index=0, raw_data=train_evaluate(baseline_parameters))
Out:

[INFO 02-03 18:40:14] ax.core.experiment: Attached custom parameterizations [{'lr': 2.6e-05, 'momentum': 0.58}] as trial 0.

Out:

[INFO 02-03 18:40:21] ax.service.ax_client: Completed trial 0 with data: {'accuracy': (0.841833, None)}.

Now we start the optimization loop.

At each step, the user queries the client for a new trial then submits the evaluation of that trial back to the client.

Note that Ax auto-selects an appropriate optimization algorithm based on the search space. For more advanced use cases that require a specific optimization algorithm, pass a generation_strategy argument into the AxClient constructor. Note that when Bayesian Optimization is used, generating new trials may take a few minutes.

for i in range(25):
parameters, trial_index = ax_client.get_next_trial()
# Local evaluation here can be replaced with deployment to external system.
ax_client.complete_trial(trial_index=trial_index, raw_data=train_evaluate(parameters))
Out:

/home/runner/work/Ax/Ax/ax/modelbridge/cross_validation.py:439: UserWarning:

Encountered exception in computing model fit quality: RandomModelBridge does not support prediction.

[INFO 02-03 18:40:22] ax.service.ax_client: Generated new trial 1 with parameters {'lr': 0.009955, 'momentum': 0.633423} using model Sobol.

Out:

[INFO 02-03 18:40:28] ax.service.ax_client: Completed trial 1 with data: {'accuracy': (0.100333, None)}.

Out:

/home/runner/work/Ax/Ax/ax/modelbridge/cross_validation.py:439: UserWarning:

Encountered exception in computing model fit quality: RandomModelBridge does not support prediction.

[INFO 02-03 18:40:28] ax.service.ax_client: Generated new trial 2 with parameters {'lr': 9e-05, 'momentum': 0.335441} using model Sobol.

Out:

[INFO 02-03 18:40:35] ax.service.ax_client: Completed trial 2 with data: {'accuracy': (0.884333, None)}.

Out:

/home/runner/work/Ax/Ax/ax/modelbridge/cross_validation.py:439: UserWarning:

Encountered exception in computing model fit quality: RandomModelBridge does not support prediction.

[INFO 02-03 18:40:35] ax.service.ax_client: Generated new trial 3 with parameters {'lr': 2e-06, 'momentum': 0.902833} using model Sobol.

Out:

[INFO 02-03 18:40:42] ax.service.ax_client: Completed trial 3 with data: {'accuracy': (0.678667, None)}.

Out:

/home/runner/work/Ax/Ax/ax/modelbridge/cross_validation.py:439: UserWarning:

Encountered exception in computing model fit quality: RandomModelBridge does not support prediction.

[INFO 02-03 18:40:42] ax.service.ax_client: Generated new trial 4 with parameters {'lr': 0.10174, 'momentum': 0.065787} using model Sobol.

Out:

[INFO 02-03 18:40:49] ax.service.ax_client: Completed trial 4 with data: {'accuracy': (0.100333, None)}.

Out:

/home/runner/work/Ax/Ax/ax/modelbridge/cross_validation.py:439: UserWarning:

Encountered exception in computing model fit quality: RandomModelBridge does not support prediction.

[INFO 02-03 18:40:49] ax.service.ax_client: Generated new trial 5 with parameters {'lr': 0.017004, 'momentum': 0.802601} using model Sobol.

Out:

[INFO 02-03 18:40:55] ax.service.ax_client: Completed trial 5 with data: {'accuracy': (0.088167, None)}.

Out:

[INFO 02-03 18:40:56] ax.service.ax_client: Generated new trial 6 with parameters {'lr': 3.5e-05, 'momentum': 0.0} using model BoTorch.

Out:

[INFO 02-03 18:41:03] ax.service.ax_client: Completed trial 6 with data: {'accuracy': (0.774, None)}.

Out:

[INFO 02-03 18:41:03] ax.service.ax_client: Generated new trial 7 with parameters {'lr': 8.3e-05, 'momentum': 1.0} using model BoTorch.

Out:

[INFO 02-03 18:41:10] ax.service.ax_client: Completed trial 7 with data: {'accuracy': (0.9055, None)}.

Out:

[INFO 02-03 18:41:11] ax.service.ax_client: Generated new trial 8 with parameters {'lr': 4.8e-05, 'momentum': 1.0} using model BoTorch.

Out:

[INFO 02-03 18:41:18] ax.service.ax_client: Completed trial 8 with data: {'accuracy': (0.8565, None)}.

Out:

[INFO 02-03 18:41:18] ax.service.ax_client: Generated new trial 9 with parameters {'lr': 0.00015, 'momentum': 1.0} using model BoTorch.

Out:

[INFO 02-03 18:41:25] ax.service.ax_client: Completed trial 9 with data: {'accuracy': (0.486167, None)}.

Out:

[INFO 02-03 18:41:26] ax.service.ax_client: Generated new trial 10 with parameters {'lr': 1.3e-05, 'momentum': 1.0} using model BoTorch.

Out:

[INFO 02-03 18:41:33] ax.service.ax_client: Completed trial 10 with data: {'accuracy': (0.8575, None)}.

Out:

[INFO 02-03 18:41:33] ax.service.ax_client: Generated new trial 11 with parameters {'lr': 1.1e-05, 'momentum': 0.0} using model BoTorch.

Out:

[INFO 02-03 18:41:40] ax.service.ax_client: Completed trial 11 with data: {'accuracy': (0.622667, None)}.

Out:

[INFO 02-03 18:41:41] ax.service.ax_client: Generated new trial 12 with parameters {'lr': 7e-05, 'momentum': 0.253302} using model BoTorch.

Out:

[INFO 02-03 18:41:48] ax.service.ax_client: Completed trial 12 with data: {'accuracy': (0.887167, None)}.

Out:

[INFO 02-03 18:41:48] ax.service.ax_client: Generated new trial 13 with parameters {'lr': 1e-06, 'momentum': 0.0} using model BoTorch.

Out:

[INFO 02-03 18:41:55] ax.service.ax_client: Completed trial 13 with data: {'accuracy': (0.133333, None)}.

Out:

[INFO 02-03 18:41:56] ax.service.ax_client: Generated new trial 14 with parameters {'lr': 0.000139, 'momentum': 0.0} using model BoTorch.

Out:

[INFO 02-03 18:42:03] ax.service.ax_client: Completed trial 14 with data: {'accuracy': (0.849833, None)}.

Out:

[INFO 02-03 18:42:04] ax.service.ax_client: Generated new trial 15 with parameters {'lr': 5e-06, 'momentum': 1.0} using model BoTorch.

Out:

[INFO 02-03 18:42:11] ax.service.ax_client: Completed trial 15 with data: {'accuracy': (0.828167, None)}.

Out:

[INFO 02-03 18:42:12] ax.service.ax_client: Generated new trial 16 with parameters {'lr': 1.9e-05, 'momentum': 1.0} using model BoTorch.

Out:

[INFO 02-03 18:42:18] ax.service.ax_client: Completed trial 16 with data: {'accuracy': (0.827, None)}.

Out:

[INFO 02-03 18:42:20] ax.service.ax_client: Generated new trial 17 with parameters {'lr': 3.9e-05, 'momentum': 0.566309} using model BoTorch.

Out:

[INFO 02-03 18:42:26] ax.service.ax_client: Completed trial 17 with data: {'accuracy': (0.874, None)}.

Out:

[INFO 02-03 18:42:27] ax.service.ax_client: Generated new trial 18 with parameters {'lr': 7.5e-05, 'momentum': 0.0} using model BoTorch.

Out:

[INFO 02-03 18:42:34] ax.service.ax_client: Completed trial 18 with data: {'accuracy': (0.810667, None)}.

Out:

[INFO 02-03 18:42:35] ax.service.ax_client: Generated new trial 19 with parameters {'lr': 2.4e-05, 'momentum': 0.763182} using model BoTorch.

Out:

[INFO 02-03 18:42:42] ax.service.ax_client: Completed trial 19 with data: {'accuracy': (0.883667, None)}.

Out:

[INFO 02-03 18:42:43] ax.service.ax_client: Generated new trial 20 with parameters {'lr': 0.000286, 'momentum': 0.0} using model BoTorch.

Out:

[INFO 02-03 18:42:50] ax.service.ax_client: Completed trial 20 with data: {'accuracy': (0.8355, None)}.

Out:

[INFO 02-03 18:42:51] ax.service.ax_client: Generated new trial 21 with parameters {'lr': 5.7e-05, 'momentum': 0.448405} using model BoTorch.

Out:

[INFO 02-03 18:42:58] ax.service.ax_client: Completed trial 21 with data: {'accuracy': (0.881, None)}.

Out:

[INFO 02-03 18:42:59] ax.service.ax_client: Generated new trial 22 with parameters {'lr': 2.8e-05, 'momentum': 0.805158} using model BoTorch.

Out:

[INFO 02-03 18:43:06] ax.service.ax_client: Completed trial 22 with data: {'accuracy': (0.888167, None)}.

Out:

[INFO 02-03 18:43:07] ax.service.ax_client: Generated new trial 23 with parameters {'lr': 0.000838, 'momentum': 0.0} using model BoTorch.

Out:

[INFO 02-03 18:43:13] ax.service.ax_client: Completed trial 23 with data: {'accuracy': (0.405333, None)}.

Out:

[INFO 02-03 18:43:15] ax.service.ax_client: Generated new trial 24 with parameters {'lr': 5.1e-05, 'momentum': 0.337101} using model BoTorch.

Out:

[INFO 02-03 18:43:21] ax.service.ax_client: Completed trial 24 with data: {'accuracy': (0.873667, None)}.

Out:

[INFO 02-03 18:43:22] ax.service.ax_client: Generated new trial 25 with parameters {'lr': 6.8e-05, 'momentum': 0.740126} using model BoTorch.

Out:

[INFO 02-03 18:43:29] ax.service.ax_client: Completed trial 25 with data: {'accuracy': (0.897833, None)}.

How many trials can run in parallel?

By default, Ax restricts number of trials that can run in parallel for some optimization stages, in order to improve the optimization performance and reduce the number of trials that the optimization will require. To check the maximum parallelism for each optimization stage:

ax_client.get_max_parallelism()
Out:

[(5, 5), (-1, 3)]

The output of this function is a list of tuples of form (number of trials, max parallelism), so the example above means "the max parallelism is 5 for the first 5 trials and 3 for all subsequent trials." This is because the first 5 trials are produced quasi-randomly and can all be evaluated at once, and subsequent trials are produced via Bayesian optimization, which converges on optimal point in fewer trials when parallelism is limited. MaxParallelismReachedException indicates that the parallelism limit has been reached –– refer to the 'Service API Exceptions Meaning and Handling' section at the end of the tutorial for handling.

How to view all existing trials during optimization?

ax_client.get_trials_data_frame()
trial_indexarm_nametrial_statusgeneration_methodgeneration_nodeaccuracylrmomentum
000_0COMPLETEDnannan0.8418332.6e-050.58
111_0COMPLETEDSobolGenerationStep_00.1003330.0099550.633423
222_0COMPLETEDSobolGenerationStep_00.8843339e-050.335441
333_0COMPLETEDSobolGenerationStep_00.6786672e-060.902833
444_0COMPLETEDSobolGenerationStep_00.1003330.101740.065787
555_0COMPLETEDSobolGenerationStep_00.0881670.0170040.802601
666_0COMPLETEDBoTorchGenerationStep_10.7743.5e-050
777_0COMPLETEDBoTorchGenerationStep_10.90558.3e-051
888_0COMPLETEDBoTorchGenerationStep_10.85654.8e-051
999_0COMPLETEDBoTorchGenerationStep_10.4861670.000151
101010_0COMPLETEDBoTorchGenerationStep_10.85751.3e-051
111111_0COMPLETEDBoTorchGenerationStep_10.6226671.1e-050
121212_0COMPLETEDBoTorchGenerationStep_10.8871677e-050.253302
131313_0COMPLETEDBoTorchGenerationStep_10.1333331e-060
141414_0COMPLETEDBoTorchGenerationStep_10.8498330.0001390
151515_0COMPLETEDBoTorchGenerationStep_10.8281675e-061
161616_0COMPLETEDBoTorchGenerationStep_10.8271.9e-051
171717_0COMPLETEDBoTorchGenerationStep_10.8743.9e-050.566309
181818_0COMPLETEDBoTorchGenerationStep_10.8106677.5e-050
191919_0COMPLETEDBoTorchGenerationStep_10.8836672.4e-050.763182
202020_0COMPLETEDBoTorchGenerationStep_10.83550.0002860
212121_0COMPLETEDBoTorchGenerationStep_10.8815.7e-050.448405
222222_0COMPLETEDBoTorchGenerationStep_10.8881672.8e-050.805158
232323_0COMPLETEDBoTorchGenerationStep_10.4053330.0008380
242424_0COMPLETEDBoTorchGenerationStep_10.8736675.1e-050.337101
252525_0COMPLETEDBoTorchGenerationStep_10.8978336.8e-050.740126

6. Retrieve best parameters

Once it's complete, we can access the best parameters found, as well as the corresponding metric values. Note that these parameters may not necessarily be the set that yielded the highest observed accuracy because Ax uses the highest model predicted accuracy to choose the best parameters (see here for more details). Due to randomness in the data or the algorithm itself, using observed accuracy may result in choosing an outlier for the best set of parameters. Using the model predicted best will use the model to regularize the observations and reduce the likelihood of picking some outlier in the data.

best_parameters, values = ax_client.get_best_parameters()
best_parameters
Out:

{'lr': 6.844706951213188e-05, 'momentum': 0.7401263354814377}

mean, covariance = values
mean
Out:

{'accuracy': 0.9022904560889984}

7. Plot the response surface and optimization trace

Contour plot showing classification accuracy as a function of the two hyperparameters.

The black squares show points that we have actually run; notice how they are clustered in the optimal region.

render(ax_client.get_contour_plot(param_x="lr", param_y="momentum", metric_name="accuracy"))
Out:

[INFO 02-03 18:43:30] ax.service.ax_client: Retrieving contour plot with parameter 'lr' on X-axis and 'momentum' on Y-axis, for metric 'accuracy'. Remaining parameters are affixed to the middle of their range.

loading...

Here we plot the optimization trace, showing the progression of finding the point with the optimal objective:

render(
ax_client.get_optimization_trace()
)
loading...

8. Train CNN with best hyperparameters and evaluate on test set

Note that the resulting accuracy on the test set generally won't be the same as the maximum accuracy achieved on the evaluation set throughout optimization.

df = ax_client.get_trials_data_frame()
best_arm_idx = df.trial_index[df["accuracy"] == df["accuracy"].max()].values[0]
best_arm = ax_client.get_trial_parameters(best_arm_idx)
best_arm
Out:

{'lr': 8.29759076064034e-05, 'momentum': 1.0}

combined_train_valid_set = torch.utils.data.ConcatDataset(
[
train_loader.dataset.dataset,
valid_loader.dataset.dataset,
]
)
combined_train_valid_loader = torch.utils.data.DataLoader(
combined_train_valid_set,
batch_size=BATCH_SIZE,
shuffle=True,
)
net = train(
net=CNN(),
train_loader=combined_train_valid_loader,
parameters=best_arm,
dtype=dtype,
device=device,
)
test_accuracy = evaluate(
net=net,
data_loader=test_loader,
dtype=dtype,
device=device,
)
print(f"Classification Accuracy (test set): {round(test_accuracy*100, 2)}%")
Out:

Classification Accuracy (test set): 39.1%

9. Save / reload optimization to JSON / SQL

We can serialize the state of optimization to JSON and save it to a .json file or save it to the SQL backend. For the former:

ax_client.save_to_json_file()  # For custom filepath, pass `filepath` argument.
Out:

[INFO 02-03 18:44:02] ax.service.ax_client: Saved JSON-serialized state of optimization to ax_client_snapshot.json.

restored_ax_client = (
AxClient.load_from_json_file()
) # For custom filepath, pass `filepath` argument.
Out:

[INFO 02-03 18:44:02] 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.

To store state of optimization to an SQL backend, first follow setup instructions on Ax website.

Having set up the SQL backend, pass DBSettings to AxClient on instantiation (note that SQLAlchemy dependency will have to be installed – for installation, refer to optional dependencies on Ax website):

from ax.storage.sqa_store.structs import DBSettings

# URL is of the form "dialect+driver://username:password@host:port/database".
db_settings = DBSettings(url="sqlite:///foo.db")
# Instead of URL, can provide a `creator function`; can specify custom encoders/decoders if necessary.
new_ax = AxClient(db_settings=db_settings)
Out:

[INFO 02-03 18:44:03] 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.

When valid DBSettings are passed into AxClient, a unique experiment name is a required argument (name) to ax_client.create_experiment. The state of the optimization is auto-saved any time it changes (i.e. a new trial is added or completed, etc).

To reload an optimization state later, instantiate AxClient with the same DBSettings and use ax_client.load_experiment_from_database(experiment_name="my_experiment").

Special Cases

Evaluation failure: should any optimization iterations fail during evaluation, log_trial_failure will ensure that the same trial is not proposed again.

_, trial_index = ax_client.get_next_trial()
ax_client.log_trial_failure(trial_index=trial_index)
Out:

[INFO 02-03 18:44:05] ax.service.ax_client: Generated new trial 26 with parameters {'lr': 8e-06, 'momentum': 0.692368} using model BoTorch.

Out:

[INFO 02-03 18:44:05] ax.service.ax_client: Registered failure of trial 26.

Need to run many trials in parallel: for optimal results and optimization efficiency, we strongly recommend sequential optimization (generating a few trials, then waiting for them to be completed with evaluation data). However, if your use case needs to dispatch many trials in parallel before they are updated with data and you are running into the "All trials for current model have been generated, but not enough data has been observed to fit next model" error, instantiate AxClient as AxClient(enforce_sequential_optimization=False).

Service API Exceptions Meaning and Handling

DataRequiredError: Ax generation strategy needs to be updated with more data to proceed to the next optimization model. When the optimization moves from initialization stage to the Bayesian optimization stage, the underlying BayesOpt model needs sufficient data to train. For optimal results and optimization efficiency (finding the optimal point in the least number of trials), we recommend sequential optimization (generating a few trials, then waiting for them to be completed with evaluation data). Therefore, the correct way to handle this exception is to wait until more trial evaluations complete and log their data via ax_client.complete_trial(...).

However, if there is strong need to generate more trials before more data is available, instantiate AxClient as AxClient(enforce_sequential_optimization=False). With this setting, as many trials will be generated from the initialization stage as requested, and the optimization will move to the BayesOpt stage whenever enough trials are completed.

MaxParallelismReachedException: generation strategy restricts the number of trials that can be run simultaneously (to encourage sequential optimization), and the parallelism limit has been reached. The correct way to handle this exception is the same as DataRequiredError – to wait until more trial evluations complete and log their data via ax_client.complete_trial(...).

In some cases higher parallelism is important, so enforce_sequential_optimization=False kwarg to AxClient allows the user to suppress limiting of parallelism. It's also possible to override the default parallelism setting for all stages of the optimization by passing choose_generation_strategy_kwargs to ax_client.create_experiment:

ax_client = AxClient()
ax_client.create_experiment(
parameters=[
{"name": "x", "type": "range", "bounds": [-5.0, 10.0]},
{"name": "y", "type": "range", "bounds": [0.0, 15.0]},
],
# Sets max parallelism to 10 for all steps of the generation strategy.
choose_generation_strategy_kwargs={"max_parallelism_override": 10},
)
Out:

[INFO 02-03 18:44:06] 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.

Out:

[INFO 02-03 18:44:06] ax.service.utils.instantiation: Inferred value type of ParameterType.FLOAT for parameter x. If that is not the expected value type, you can explicitly specify 'value_type' ('int', 'float', 'bool' or 'str') in parameter dict.

Out:

[INFO 02-03 18:44:06] ax.service.utils.instantiation: Inferred value type of ParameterType.FLOAT for parameter y. If that is not the expected value type, you can explicitly specify 'value_type' ('int', 'float', 'bool' or 'str') in parameter dict.

Out:

[INFO 02-03 18:44:06] ax.service.utils.instantiation: Created search space: SearchSpace(parameters=[RangeParameter(name='x', parameter_type=FLOAT, range=[-5.0, 10.0]), RangeParameter(name='y', parameter_type=FLOAT, range=[0.0, 15.0])], parameter_constraints=[]).

Out:

[INFO 02-03 18:44:06] ax.modelbridge.dispatch_utils: Using Models.BOTORCH_MODULAR since there is at least one ordered parameter and there are no unordered categorical parameters.

Out:

[INFO 02-03 18:44:06] ax.modelbridge.dispatch_utils: Calculating the number of remaining initialization trials based on num_initialization_trials=None max_initialization_trials=None num_tunable_parameters=2 num_trials=None use_batch_trials=False

Out:

[INFO 02-03 18:44:06] ax.modelbridge.dispatch_utils: calculated num_initialization_trials=5

Out:

[INFO 02-03 18:44:06] ax.modelbridge.dispatch_utils: num_completed_initialization_trials=0 num_remaining_initialization_trials=5

Out:

[INFO 02-03 18:44:06] ax.modelbridge.dispatch_utils: verbose, disable_progbar, and jit_compile are not yet supported when using choose_generation_strategy with ModularBoTorchModel, dropping these arguments.

Out:

[INFO 02-03 18:44:06] ax.modelbridge.dispatch_utils: Using Bayesian Optimization generation strategy: GenerationStrategy(name='Sobol+BoTorch', steps=[Sobol for 5 trials, BoTorch for subsequent trials]). Iterations after 5 will take longer to generate due to model-fitting.

ax_client.get_max_parallelism()  # Max parallelism is now 10 for all stages of the optimization.
Out:

[(5, 10), (-1, 10)]