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, CNN
init_notebook_plotting()
[INFO 01-18 03:39:10] ipy_plotting: Injecting Plotly library into cell. Do not overwrite or delete cell.
torch.manual_seed(12345)
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.
BATCH_SIZE = 512
train_loader, valid_loader, test_loader = load_mnist(batch_size=BATCH_SIZE)
1.1%
Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz to ./data/MNIST/raw/train-images-idx3-ubyte.gz
100.1%
Extracting ./data/MNIST/raw/train-images-idx3-ubyte.gz to ./data/MNIST/raw
2.0%
Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz to ./data/MNIST/raw/train-labels-idx1-ubyte.gz Extracting ./data/MNIST/raw/train-labels-idx1-ubyte.gz to ./data/MNIST/raw Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz to ./data/MNIST/raw/t10k-images-idx3-ubyte.gz
180.4%
Extracting ./data/MNIST/raw/t10k-images-idx3-ubyte.gz to ./data/MNIST/raw Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz to ./data/MNIST/raw/t10k-labels-idx1-ubyte.gz Extracting ./data/MNIST/raw/t10k-labels-idx1-ubyte.gz to ./data/MNIST/raw 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 = 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,
)
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 01-18 03:39:13] ax.modelbridge.dispatch_utils: Using Bayesian Optimization generation strategy: GenerationStrategy(name='Sobol+GPEI', steps=[Sobol for 5 arms, GPEI for subsequent arms], generated 0 arm(s) so far). Iterations after 5 will take longer to generate due to model-fitting. [INFO 01-18 03:39:13] ax.service.managed_loop: Started full optimization with 20 steps. [INFO 01-18 03:39:13] ax.service.managed_loop: Running optimization trial 1... [INFO 01-18 03:39:23] ax.service.managed_loop: Running optimization trial 2... [INFO 01-18 03:39:32] ax.service.managed_loop: Running optimization trial 3... [INFO 01-18 03:39:42] ax.service.managed_loop: Running optimization trial 4... [INFO 01-18 03:39:51] ax.service.managed_loop: Running optimization trial 5... [INFO 01-18 03:40:00] ax.service.managed_loop: Running optimization trial 6... [INFO 01-18 03:40:11] ax.service.managed_loop: Running optimization trial 7... [INFO 01-18 03:40:21] ax.service.managed_loop: Running optimization trial 8... [INFO 01-18 03:40:32] ax.service.managed_loop: Running optimization trial 9... [INFO 01-18 03:40:42] ax.service.managed_loop: Running optimization trial 10... [INFO 01-18 03:40:53] ax.service.managed_loop: Running optimization trial 11... [INFO 01-18 03:41:04] ax.service.managed_loop: Running optimization trial 12... [INFO 01-18 03:41:14] ax.service.managed_loop: Running optimization trial 13... [INFO 01-18 03:41:25] ax.service.managed_loop: Running optimization trial 14... [INFO 01-18 03:41:36] ax.service.managed_loop: Running optimization trial 15... [INFO 01-18 03:41:46] ax.service.managed_loop: Running optimization trial 16... [INFO 01-18 03:41:57] ax.service.managed_loop: Running optimization trial 17... [INFO 01-18 03:42:07] ax.service.managed_loop: Running optimization trial 18... [INFO 01-18 03:42:17] ax.service.managed_loop: Running optimization trial 19... [INFO 01-18 03:42:27] ax.service.managed_loop: Running optimization trial 20...
We can introspect the optimal parameters and their outcomes:
best_parameters
{'lr': 9.791037646113534e-05, 'momentum': 0.8132115442490506}
means, covariances = values
means, covariances
({'accuracy': 0.9271593935016549},
{'accuracy': {'accuracy': 1.296505675484216e-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='11_0', parameters={'lr': 9.791037646113534e-05, 'momentum': 0.8132115442490506})
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.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.8%
Total runtime of script: 4 minutes, 12.78 seconds.