Lavanya Shukla, ML engineer at Weights & Biases

Hyperparameter Tuning for Keras and Pytorch models

We’re excited to launch a powerful and efficient way to do hyperparameter tuning and optimization - W&B Sweeps, in both Keras and Pytoch.

With just a few lines of code Sweeps automatically search through high dimensional hyperparameter spaces to find the best performing model, with very little effort on your part.

Here’s how you can launch sophisticated hyperparameter sweeps in 3 simple steps.

Try Sweeps in Colab →

0. Integrate W&B

First let’s install the Weights & Biases library and add it into your training script.

A. Install wandb

pip install wandb
wandb login

B. Add W&B to your training script

At the top:
To your .fit() function
# train.py
import wandb
from wandb.keras import WandbCallback

# define model architecture
# compile the model
model.compile(loss='categorical_crossentropy', optimizer=sgd, metrics=['accuracy'])

# add the WandbCallback()
model.fit(X_train, y_train,  validation_data=(X_test, y_test), epochs=config.epochs,
   callbacks=[WandbCallback(data_type="image", labels=labels)])


1. Define the Sweep

You can define powerful sweeps simply by creating a YAML file that specifies the parameters to search through, the search strategy, and the optimization metric.

Here’s an example:

# sweep.yaml
program: train.py
method: random
 name: val_loss
 goal: minimize
   min: 0.00001
   max: 0.1
   values: ["adam", "sgd"]
   values: [96, 128, 148]
   value: 27
   type: hyperband
   s: 2
   eta: 3
   max_iter: 27

Let’s break this yaml file down:

You can find a list of all the configuration options here.

2. Setup a new sweep

Run wandb sweep  with the config file you created in step 1.

This creates your sweep, and returns both a unique identifier (SWEEP_ID) and a URL to track all your runs.

wandb sweep sweep.yaml

3. Launch the sweep

It’s time to launch our sweep and train some models!

You can do so by calling wandb agent with the SWEEP_ID you got from step 2.

wandb agent SWEEP_ID

This will start training models with different hyperparameter combinations and return a URL where you can track the sweep’s progress. You can launch multiple agents concurrently. Each of these agents will fetch parameters from the W&B server and use them to train the next model.

And voila! That's all there is to running a hyperparameter sweep!

Let’s see how we can extract insights about our model from sweeps next.

4. Visualize Sweep Results

Parallel coordinates plot

This plot maps hyperparameter values to model metrics. It’s useful for honing in on combinations of hyperparameters that led to the best model performance.

Hyperparameter Importance Plot

The hyperparameter importance plot surfaces which hyperparameters were the best predictors of, and highly correlated to desirable values for your metrics.

These visualizations can help you save both time and resources running expensive hyperparameter optimizations by honing in on the parameters (and value ranges) that are the most important, and thereby worthy of further exploration.

Next step - Get your hands dirty with sweeps

We created a simple training script and a few flavors of sweep configs for you to play with. We highly encourage you to give these a try. This repo also has examples to help you try more advanced sweep features like Bayesian Hyperband, and Hyperopt.

More Resources

Join our mailing list to get the latest machine learning updates.