W&B 💘 Keras

Use Weights & Biases for machine learning experiment tracking, model checkpointing, and project collaboration.

Keras Documentation

For a full guide of how to log to Weights & Biases using Keras, see our Keras documentation

What this notebook covers:

We show you how to integrate Weights & Biases with your Keras code to add experiment tracking to your pipeline. That includes:

1. Storing hyperparameters and metadata in a config.
2. Passing the wandb Keras callbacks to model.fit. This will automatically log training metrics, like loss, and system metrics, like GPU and CPU utilization.
3. Using the wandb.log API to log custom metrics.

all using the CIFAR-10 dataset.

Then, we'll show you how to catch your model making mistakes by logging both the output predictions and the input images the network used to generate them.

Follow along with a video tutorial!

Note: Sections starting with Step are all you need to integrate W&B in an existing pipeline. The rest just loads data and defines a model.

Install, Import, and Log In

Step 0: Install W&B

Step 1: Import W&B and Login

Side note: If this is your first time using W&B or you are not logged in, the link that appears after running wandb.login will take you to sign-up/login page. Signing up is easy!

Download and Prepare the Dataset

Define the Model

Here, we define a standard CNN (with convolution and max-pooling) in Keras.

Train the Model

Step 2: Give wandb.init your config

You first initialize your wandb run, letting us know some training is about to happen. Check out the official documentation for .init here $\rightarrow$

That's when you need to set your hyperparameters. They're passed in as a dictionary via the config argument, and then become available as the config attribute of wandb.

Learn more about config in this Colab Notebook $\rightarrow$

Step 3: Pass WandbMetricsLogger and WandbModelCheckpoint to model.fit

Keras has a robust callbacks system that allows users to separate model definition and the core training logic from other behaviors that occur during training and testing.

That includes, for example,

Click on the Project page link above to see your results!

Use wandb.log for custom metrics

Here, we log the error rate on the test set.

Log predictions on test data using WandbEvalCallback

The WandbEvalCallback is an abstract base class to build Keras callbacks for primarily model prediction visualization and secondarily dataset visualization.

This is a dataset and task agnostic abstract callback. To use this, inherit from this base callback class and implement the add_ground_truth and add_model_prediction methods.

The WandbEvalCallback is a utility class that provides helpful methods to:

• create data and prediction wandb.Table instances,
• log data and prediction Tables as wandb.Artifact,
• logs the data table on_train_begin,
• logs the prediction table on_epoch_end.

As an example, we have implemented WandbClsEvalCallback below for an image classification task. This example callback:

• logs the validation data (data_table) to W&B,
• performs inference and logs the prediction (pred_table) to W&B on every epoch end.

How the memory footprint is reduced?

We log the data_table to W&B when the on_train_begin method is ivoked. Once it's uploaded as a W&B Artifact, we get a reference to this table which can be accessed using data_table_ref class variable. The data_table_ref is a 2D list that can be indexed like self.data_table_ref[idx][n] where idx is the row number while n is the column number. Let's see the usage in the example below.

Sub-class WandbEvalCallback

Create Dataset processing and Dataloaders functions

Define our model

Set our config

Dataset

In this example, we will be using CIFAR100 dataset from TensorFlow Dataset catalog. We aim to build a simple image classification pipeline using TensorFlow/Keras.

Create our Dataloaders and Model

Train the model and log the predictions to a W&B Table

Click on the W&B project page link above to see your live results.

Whats Next? Hyperparameters with Sweeps

We tried out two different hyperparameter settings by hand. You can use Weights & Biases Sweeps to automate hyperparameter testing and explore the space of possible models and optimization strategies.

Check out Hyperparameter Optimization in TensorFlow uisng W&B Sweep $\rightarrow$

Running a hyperparameter sweep with Weights & Biases is very easy. There are just 3 simple steps:

1. Define the sweep: We do this by creating a dictionary or a YAML file that specifies the parameters to search through, the search strategy, the optimization metric et all.

2. Initialize the sweep: sweep_id = wandb.sweep(sweep_config)

3. Run the sweep agent: wandb.agent(sweep_id, function=train)

And voila! That's all there is to running a hyperparameter sweep! In the notebook below, we'll walk through these 3 steps in more detail.

Documentation

For a full guide of how to log to Weights & Biases using Keras, see our Keras documentation