Skip to main content

Optimizing CIFAR-10 Hyperparameters with W&B and SageMaker

Hyperparameter tuning for CIFAR-10 is easier than ever with AWS SageMaker and Weights & Biases, helping you find the best model settings with minimal effort.
Created on September 21|Last edited on February 8
Everyone knows that hyperparameter sweeps are a great way to get an extra level of performance out of your algorithm, but we often don’t do them because they’re expensive and tricky to set up. AWS SageMaker simplifies the process by handling hyperparameter tuning at scale, while Weights & Biases makes it effortless to track and visualize the results.

The scenario

I had code for a convolutional neural network (CNN) to classify images in the CIFAR-10 dataset, and I wanted to find the best set of hyperparameters. Instead of manually tuning parameters, I used SageMaker to automate the process and Weights & Biases to track and analyze the results.


Setting up hyperparameter tuning in SageMaker

AWS SageMaker automatically runs multiple training jobs with different hyperparameter values to find the best-performing configuration. Here’s how I set up a hyperparameter sweep using SageMaker’s HyperparameterTuner API.

Defining the estimator

First, I created a SageMaker estimator to define how the model should be trained:
estimator = PyTorch(entry_point="cifar10.py",
source_dir=os.getcwd() + "/source",
role=role,
framework_version='1.0.0.dev',
train_instance_count=1,
train_instance_type='ml.c5.xlarge',
hyperparameters={
'epochs': 50,
'momentum': 0.9
})

hyperparameter_ranges = {
'lr': ContinuousParameter(0.0001, 0.001),
'hidden_nodes': IntegerParameter(20, 100),
'batch_size': CategoricalParameter([128, 256, 512]),
'conv1_channels': CategoricalParameter([32, 64, 128]),
'conv2_channels': CategoricalParameter([64, 128, 256, 512]),
}
SageMaker automatically selects different combinations of these hyperparameters, launching training jobs in parallel to identify the optimal set.

Tracking and visualizing results with Weights & Biases
Each training job produces different results depending on the chosen hyperparameters. Weights & Biases helps monitor and compare these results with interactive visualizations.

Logging sweep results
To track the results, I integrated Weights & Biases by logging hyperparameters and metrics:
SageMaker automatically selects different combinations of these hyperparameters, launching training jobs in parallel to identify the optimal set.

Tracking and visualizing results with Weights & Biases

Each training job produces different results depending on the chosen hyperparameters. Weights & Biases helps monitor and compare these results with interactive visualizations.

Logging sweep results

To track the results, I integrated Weights & Biases by logging hyperparameters and metrics. Here’s a W&B table where I’m tracking all the runs that ran in the sweep, sorted by test accuracy. The test accuracy ranges from 10 - 76.45%, depending on the hyperparameters.

To dig deeper into the patterns between hyperparameters and accuracy on different classes, we generated a parallel coordinates plot in Weights & Biases. The first five columns are the configuration parameters and the far right column is the test accuracy.
Each line corresponds to a single run. I’ve highlighted the runs with the best test accuracy to see where they land on the other columns. I discovered that lower learning rate (config:lr) and fewer hidden nodes (config:hidden_nodes) correlated with higher test accuracy.

We can zoom in on individual accuracy metrics or even compare our different models’ classification on a single image. Here are the top 10 models overall, struggling to identify this picture of a dog correctly.


Conclusion

By leveraging AWS SageMaker for automated hyperparameter tuning and Weights & Biases for tracking and visualization, I was able to efficiently optimize my CNN for CIFAR-10 classification.
If you want to try this setup yourself, you can find the complete code on GitHub. For more details on integrating Weights & Biases with SageMaker, check out our official documentation.


Iterate on AI agents and models faster. Try Weights & Biases today.