🍀 Integrate Weights and Biases in your TensorFlow/Keras workflow

In this example, we will we will train an image classifier for the bloodMNIST dataset. The primary focus however will be the use of Weights and Biases and how easily it can be added in your TensorFlow/Keras workflow.

Consider Weights and Biases (W&B) to be the GitHub for machine learning. Use W&B for machine learning experiment tracking, dataset and model versioning, project collaboration, hyperparameter optimization, dataset exploration, model evaluation and so much more.

W&B comes with a lightweight integration for Keras (WandbCallback) and with just a few lines of code you can log your metrics, save model, training configuration, evaluate model and more. W&B is intrumented with most of your favourite machine learning frameworks.

This notebook also introduces W&B Tables. Tables accelerate the ML development lifecycle by giving users the ability to rapidly extract meaningful insights from the data. The WB Table Visualizer provides an interactive interface to perform powerful analytics functions like grouping, joining, and creating custom fields while simultaneously supporting rich media annotations such as bounding boxes and segmentation masks.

🌴 In this Notebook

In this colab we'll cover:

• training an image classifier for medMNIST (bloodMNIST) dataset,
• use of W&B Tables for dataset exploration and evaluation,
• WandbCallback for experiment tracking and model evaluation.

In addition we will also cover some of the best practices of using Weights and Biases to get the most out of your data and model.

We will start by installing the dependencies and importing required libraries.

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 for a free account is as easy as a few clicks.

🎋 Configs

Configuration files in .yaml or .json format is an integral part of most mature machine learning systems. Keeping the track of hyperparameters used to train/evaluate your model is essential for reproducing the experiments.

W&B can keep track of your configs. Here we will first define all the hyperparameters needed for training our classifier.

🍁 Prepare Dataset

MedMNIST is a large-scale MNIST-like collection of standardized biomedical images, including 12 datasets for 2D and 6 datasets for 3D. All the images are pre-processed to image size of 28x 28 and doesn't require any prior domain knowledge to start with.

In this tutorial, we will be using BloodMNIST dataset. From the dataset description.

The BloodMNIST is based on a dataset of individual normal cells, captured from individuals without infection, hematologic or oncologic disease and free of any pharmacologic treatment at the moment of blood collection. It contains a total of 17,092 images and is organized into 8 classes. We split the source dataset with a ratio of 7:1:2 into training, validation and test set. The source images with resolution 3×360×363 pixels are center-cropped into 3×200×200, and then resized into 3×28×28.

Each MedMNIST dataset can be downloaded using the download_and_prepare_dataset function below and the downloaded dataset is in the .npz format.

Each subset (e.g., bloodmnist.npz) is comprised of 6 keys: train_images, train_labels, val_images, val_labels, test_images and test_labels.

🌳 Explore the Dataset using W&B Tables

As TensorFlow/Keras users you might be familar with the show_batch function. Or you might have written some matplotlib based code to visualize few batches of dataset. This is good for quick inspection of the dataset but for most real life scenario it's not enough.

Here we will use W&B Tables (wandb.Table) to log the training data and visualize and query iteractively with W&B. As the name suggests it is a table of data specified by you. Check out more on Tables here.

You can log data to W&B Tables row wise or column wise. In the section below, we have created the table column wise. Use add_column to define the name of the column and provide array of data associated with that column. Simply adding array of images will not render in the W&B Tables UI. You will have to wrap each image array with wandb.Image. To do so, add_computed_columns is used. You can learn about these methods here.

Finally, note that W&B Tables is built on top of W&B Artifacts, which can be viewed as a file (usually for dataset and models) storage system in W&B. In this section, we have explicitly initialized an artifact using wandb.Artifact and have added both the train_table and validation_table to the artifact. Alternatively, we could have prepared the table and logged it using wandb.log. Here's a quick example if you are interested.

If you want to log the entire dataset place a tick in the log_full checkbox. Note that we are logging the entire validation data.

🌲 Data Pipeline

tf.data.Dataset is used to build the data pipeline.

🪴 Data Augmentation

We will apply simple image augmentation policies using the Keras preprocessing layers API.

🌿 Visualize Different Augmented View.

Here, let's use W&B Tables to visualize augmented images of a subset of training images.

Augmentation policies should make sense for the given classification task. By using W&B Tables here we can visualize how the original images are augmented. For the sake of simplicity, we will just be visualizing the first 100 images.

We can download the dataset that we have logged as W&B Tables as shown in the code cell below. Since Tables are saved as W&B Artifacts, we first need to pass in the name (path as shown in the UI) of the artifact to use_artifact. You can find the name if you head over to the artifact tab on the W&B dashboard and click on the API panel.

Get the required table by using the get method and provide the name of the table. Use the get_column method get the data associated with that column. Here, the augment_table is initialized with the column names and data are added row-wise iteratively.

🎄 Model

We will be using VGG16 as the backbone CNN block.

☘️ Callback

Here we will define early stopping callback. We will define the WandbCallback later.

You can use wandb.log to log any useful metric/parameter that's not logged by WandbCallback. Here we are using a learning rate scheduler to exponentially decay the learning rate after 10 epochs. Notice the use of wandb.log to capture the learning rate and commit=False in particular.

You can learn more about wandb.log here.

🌻 Train

🍄 Train using WandbCallback

In the section below we will train our classifier using WandbCallback to by default log all the training and validation metrics to a wandb dashboard.

WandbCallback enables to you keep track of your experiments, saves the best model, and helps visualize model performance with just one line of code.

In the section below, we have used the following arguments:

• monitor = 'val_loss' will monitor the mentioned metric to save the best model. Note that 'val_loss' is the default value for monitor.
• log_weights = True save histograms of the model's layer's weights.
• log_evaluation = True will create a W&B Table of validation data and model prediction. The number of validation samples is controlled by validation_steps if a generator is passed to model.fit.

Check out the documentation here to know more about the WandbCallback.

🌱 Advanced Usage

In this section, we will see an advance usage of WandbCallback.

We will use WandbCallback to log the GradCAM for each validation examples along with gound truth labels and model predictions.

We will be using this tutorial on GradCAM by François Chollet.

In the cell block below, we will be using WandbCallback's validation_row_processor and prediction_row_processor to log the images, ground truth label, model prediction and the GradCAM for model interpretability.

The processors' take a callable function that receive an ndx (index) and a row (dict of data). The validation_processor function below receives the input image array along with target label as row dict. The prediction_processor receives the model output prediction and the validation data row index.

The validation_row_processor is executed when WandbCallback is initialized (i.e, before model training) while prediction_row_processor is called once the training is over. The validation_row_processor creates a table with two columns input:image and target:class. Notice that in the prediction_processor function we can get the logged image at a given val_row using the get_row method.

🌾 Conclusion

Weights and Biases's Keras integration enables experiment tracking and so much more with just few lines of code. In this notebook, we have seen some advanced usage of Keras WandbCallback and different ways of using W&B Tables for evaluation and data exploration.

To sum up all you need is a free W&B account, import the WandbCallback and pass it to model.fit(callbacks=[.]) just like any callback. There are few more arguments that you can learn about in the documentation page here. In particular,

• you can log the metrics for each batch by setting log_batch_frequency=1,
• you can log the gradients of each layer to debug vanishing or exploding gradient issue by setting log_gradients=True. You will also have to provide the training_data in the format of (X, y).
• if your task is semantic segmentation you can set input_type=segmentation_mask.

If a usecase is not covered by the WandbCallback you can easily write a custom Keras callback and use wandb.log to log the required data to W&B dashboard.