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 torch
import numpy as np
from ax.plot.contour import plot_contour
from ax.plot.trace import optimization_trace_single_method
from ax.service.managed_loop import optimize
from ax.utils.notebook.plotting import render, init_notebook_plotting
from ax.utils.tutorials.cnn_utils import load_mnist, train, evaluate
init_notebook_plotting()
[INFO 04-30 12:37:08] ipy_plotting: Injecting Plotly library into cell. Do not overwrite or delete cell.
dtype = torch.float
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
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.
train_loader, valid_loader, test_loader = load_mnist()
0%| | 0.00/9.91M [00:00<?, ?B/s]
Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz to ./data/raw/train-images-idx3-ubyte.gz
9.92MB [00:02, 4.45MB/s] 32.8KB [00:00, 197KB/s] 0.00B [00:00, ?B/s]
Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz to ./data/raw/train-labels-idx1-ubyte.gz Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz to ./data/raw/t10k-images-idx3-ubyte.gz
1.65MB [00:03, 542KB/s] 0.00B [00:00, ?B/s]
Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz to ./data/raw/t10k-labels-idx1-ubyte.gz
8.19KB [00:00, 17.1KB/s]
Processing... Done!
In this tutorial, we want to optimize classification accuracy on the validation set as a function of the learning rate and momentum. The function takes in a parameterization (set of parameter values), computes the classification accuracy, and returns a dictionary of metric name ('accuracy') to a tuple with the mean and standard error.
def train_evaluate(parameterization):
net = train(train_loader=train_loader, parameters=parameterization, dtype=dtype, device=device)
return evaluate(
net=net,
data_loader=valid_loader,
dtype=dtype,
device=device,
)
Here, we set the bounds on the learning rate and momentum and set the parameter space for the learning rate to be on a log scale.
best_parameters, values, experiment, model = optimize(
parameters=[
{"name": "lr", "type": "range", "bounds": [1e-6, 0.4], "log_scale": True},
{"name": "momentum", "type": "range", "bounds": [0.0, 1.0]},
],
evaluation_function=train_evaluate,
objective_name='accuracy',
)
[INFO 04-30 12:37:40] ax.service.utils.dispatch: Using Bayesian Optimization generation strategy. Iterations after 5 will take longer to generate due to model-fitting. [INFO 04-30 12:37:40] ax.service.managed_loop: Started full optimization with 20 steps. [INFO 04-30 12:37:40] ax.service.managed_loop: Running optimization trial 1... [INFO 04-30 12:38:49] ax.service.managed_loop: Running optimization trial 2... [INFO 04-30 12:39:57] ax.service.managed_loop: Running optimization trial 3... [INFO 04-30 12:41:05] ax.service.managed_loop: Running optimization trial 4... [INFO 04-30 12:42:13] ax.service.managed_loop: Running optimization trial 5... [INFO 04-30 12:43:20] ax.service.managed_loop: Running optimization trial 6... [INFO 04-30 12:44:42] ax.service.managed_loop: Running optimization trial 7... [INFO 04-30 12:46:06] ax.service.managed_loop: Running optimization trial 8... [INFO 04-30 12:47:34] ax.service.managed_loop: Running optimization trial 9... [INFO 04-30 12:49:00] ax.service.managed_loop: Running optimization trial 10... [INFO 04-30 12:50:24] ax.service.managed_loop: Running optimization trial 11... [INFO 04-30 12:52:01] ax.service.managed_loop: Running optimization trial 12... [INFO 04-30 12:53:32] ax.service.managed_loop: Running optimization trial 13... [INFO 04-30 12:55:08] ax.service.managed_loop: Running optimization trial 14... [INFO 04-30 12:56:45] ax.service.managed_loop: Running optimization trial 15... [INFO 04-30 12:58:17] ax.service.managed_loop: Running optimization trial 16... [INFO 04-30 12:59:55] ax.service.managed_loop: Running optimization trial 17... [INFO 04-30 13:01:43] ax.service.managed_loop: Running optimization trial 18... [INFO 04-30 13:03:48] ax.service.managed_loop: Running optimization trial 19... [INFO 04-30 13:06:00] ax.service.managed_loop: Running optimization trial 20...
We can introspect the optimal parameters and their outcomes:
best_parameters
{'lr': 0.0029176399675537317, 'momentum': 3.0347402313065844e-16}
means, covariances = values
means, covariances
({'accuracy': 0.968833362542745},
{'accuracy': {'accuracy': 1.3653840299223108e-08}})
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(plot_contour(model=model, param_x='lr', param_y='momentum', metric_name='accuracy'))
Show the model accuracy improving as we identify better hyperparameters.
# `plot_single_method` expects a 2-d array of means, because it expects to average means from multiple
# optimization runs, so we wrap out best objectives array in another array.
best_objectives = np.array([[trial.objective_mean*100 for trial in experiment.trials.values()]])
best_objective_plot = optimization_trace_single_method(
y=np.maximum.accumulate(best_objectives, axis=1),
title="Model performance vs. # of iterations",
ylabel="Classification Accuracy, %",
)
render(best_objective_plot)
Note that the resulting accuracy on the test set might not be exactly the same as the maximum accuracy achieved on the evaluation set throughout optimization.
data = experiment.fetch_data()
df = data.df
best_arm_name = df.arm_name[df['mean'] == df['mean'].max()].values[0]
best_arm = experiment.arms_by_name[best_arm_name]
best_arm
Arm(name='17_0', parameters={'lr': 0.0029176399675537317, 'momentum': 3.0347402313065844e-16})
net = train(
train_loader=train_loader,
parameters=best_arm.parameters,
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)}%")
Classification Accuracy (test set): 97.06%