How to use Wandb for Experimentation?

Wandb (Weights and Biases) is a powerful platform designed to streamline and enhance deep learning experimentation. It provides a unified interface to track, visualize, and collaborate on machine learning projects. With it, researchers and developers can effortlessly log and monitor various metrics, hyper-parameters, and system resources during training, enabling them to gain valuable insights into their models’ performance.

You can find the notebook version here.

Uses of Wandb

• Tracking and Visualization: It provides a unified interface to track and visualize metrics, hyperparameters, and system resources during deep learning experiments.
• Interactive Dashboards: The logged metrics are automatically organized and presented in intuitive and interactive dashboards. This feature makes it easy to analyze and compare different experiments.
• Hyper-parameter Tracking: It offers powerful hyperparameter tracking capabilities. You can log and compare hyperparameter settings across experiments.
• Advanced Search: With it, you can leverage advanced search capabilities to find the best hyperparameter configurations for your models. This feature helps you efficiently explore the hyperparameter space and identify optimal settings.
• Collaboration and Knowledge Sharing: It facilitates collaboration among team members by providing a centralized platform.
• Framework Integration: It seamlessly integrates with popular deep learning frameworks like TensorFlow and PyTorch, as well as other tools in the machine learning ecosystem.

About this notebook

In this notebook we are going to see some of the functionalities of the wandb for deep learning experimentation. As we know, any deep learning experimentation is computationally heavy and may lead to wrong results. To solve this problem we need tools like wandb which help us analyse our model performance as well as give us direction with which we can proceed to further experimentation.

You can check out the official documentation here.

Functionalities of Wandb

Here are the functionalities which we are going to see in this notebook,

1. 💻 Wandb config : Config helps track experiments by allowing you to easily log and track hyperparameters and their values. With Wandb config, you can define and set hyperparameters within your code, and Wandb will automatically log and organize them for each experiment. This enables you to keep track of the specific hyperparameter values used in each run, making it easier to compare and analyze their impact on your model’s performance.
2. 🔥 Wandb init : Init helps track experiments by initializing a project and connecting it to the platform. When we run wandb.init in our code, it creates a new run in your project and assigns it a unique ID. This run ID is used to identify and track the specific experiment
3. 👀 Wandb watch: Watch helps track experiments by automatically monitoring and logging the gradients and parameters of your machine learning model during training. By using the wandb.watch() function in your code, Wandb is able to keep track of these values and visualize them on the Wandb dashboard.
4. 🪵 Wandb log: log helps track experiments by allowing you to log various metrics and other relevant information during the course of your deep learning experiment. By using the wandb.log() function in your code, you can log key metrics such as loss, accuracy, and custom evaluation metrics at different stages of your training or evaluation process.
5. 💾 Wandb save: save helps track experiments by allowing you to save and log important artefacts such as model weights, trained models, datasets, and other files relevant to your deep learning experiment. By using the wandb.save()function in your code, you can specify files or directories that you want to save and associate with your experiment run.

Flow of the Notebook

We are going to build a CNN classification model on Sign Langugae Recognition MNIST dataset in this notebook. We will track the entire experimentation process with the help of Wandb.

The flow of this notebook follows the usual Pytorch deep learning experimentation.

1. First we collect the data and pre-process to our need (demo purposes)
2. Setting up the Wandb with your own account
3. Define the model hyperparameteres using config
4. Define the pipeline for the model for wandb.init
5. Train and save the model into ONNX format

Without further ado, let’s get started…

The original blog for Sign Language can be viewed here

About the Sign Language Dataset

Sign Language dataset contains American Sign Language alphabet images. There are total 24 alphabets in this dataset (‘J’ and ‘Z’ missing) as this is a image dataset and the symbols which require motion (video input) are ignored.

You can find the entire dataset here.

Importing essential libraries

First step is to import all the necessary libraries. For this notebook we are going to use Pytorch. We have imported usual stuff that we require for any deep learning classification task below,

import torch 
import torch.nn as nn 
import pandas as pd 
import numpy as np 
import matplotlib.pyplot as plt 
from sklearn.model_selection import train_test_split
from torch.utils.data import Dataset, DataLoader
from torchvision.transforms import transforms
from PIL import Image
from tqdm import tqdm

Setting up device

PyTorch requires us to define the device on which we are going to train the model.

device = 'cuda' if torch.cuda.is_available() else 'cpu'

Importing data from Kaggle using API

To import the data directly from the Kaggle we can use the following cell. You would need to insert your kaggle key into the colab before running this operation.

To understand this, you can refer this video.

# setting up kaggle API
! mkdir ~/.kaggle
! cp kaggle.json ~/.kaggle/

# Downloading the dataset directly from kaggle
! kaggle datasets download -d datamunge/sign-language-mnist

# Unzipping the data
! unzip /content/sign-language-mnist.zip

Train Test Split

There are two separate files for training and testing. Both of these files contain 784 columns which are pixel values for 28×28 pixel image.

Train data contains the label column as well.

train = pd.read_csv('/content/sign_mnist_train/sign_mnist_train.csv')
test = pd.read_csv('/content/sign_mnist_test/sign_mnist_test.csv')

Image data

X = train.drop(['label'], axis = 1) # Detaching label
y = train['label'] # selecting target

We have more than 27 thousand images which is enough to train a CNN model.

X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=42)

We can use the following code to see how the actual images are mapped.

im = np.array(X_train.iloc[2])
im = im.reshape((28,28, 1))
plt.imshow(im)

Now that we have the data ready, let’s start with our Deep learning eperimenation.

Setting up Wandb

You would need to have a Wandb account in order to use any functionalities. You can create your own account going directly to the website.

Install wandb

Installing wandb is as simple as any other python package using PIP installer. We can use the following cell to install it on colab notebook.

! pip install wandb

Login into Wandb

Once we have our Wandb account ready it is easy to access the API. We just need to login from colab using the following cell.

You would require to insert the API key in order to run this cell. You can copy your own API key from here.

! wandb login

Import Wandb

import wandb

Wandb.init: Initialize Project

Once we have everything setup, we can now initialize our wandb project using the wandb.init().

We have provided the project name and name of the instance to wandb.init. You can set up any project name.

This cell will give you the link for the dashboard where you can find the project and all the initial setup.

wandb.init(project="sign-language-model", name="SLR")

Defining Hyper-Parameters

We will define the hyperparameters for our CNN model. The config is nothing but a dictonary with all the hyper-parameters. We can set this dictionary as a config dictionary.

config = dict(
    epochs= 20,           # no of epochs
    classes=24 +1,        # classes (total 24 letters plus additional None)
    image_size = 28,      # size of the image
    kernels=[16, 32],     # kernel size for each layer in CNN, you can tweak this according to your need
    batch_size=32,
    learning_rate=0.005,
    architecture="CNN"
)

Experimental Pipeline

This is the most important step, as we are using plain Pytorch (not Pytorch lightning or Keras over Tensorflow) we need to define the moddel pipeline. Following code defines everything that we are going to do for this model training.

Let’s put it stepwise,

1. Initialize the wandb: We initialize the project (just like we did before). This will create a project with name specified and the instance of that project in our dashboard. It also takes the hypyerparameter as the input, wandb tracks all the values of the hyperparameters provided on the specified project d
2. Passing the config: We use wandb.config to track all the parameters provided in our config. This help us track the model performance based on the various parameters selected.
3. Training and testing the model

wandb.config

def model_pipeline(hyperparameters):

    # tell wandb to get started
    with wandb.init(project="sign-language-model", name="SLR", config =hyperparameters):
      # access all HPs through wandb.config, so logging matches execution!
      config = wandb.config

      # make the model, data, and optimization problem
      model, train_loader, test_loader, criterion, optimizer = model_creation(config)
      print(model)

      # and use them to train the model
      train_model(model, train_loader, criterion, optimizer, config)

      test_model(model, test_loader)

    return model
     

You can see these details in the dashboard under overview,

Dataset and DataLoaders

For this demo notebook, we are doing just basic pre-processing for our image data. This code is sufficient for most of the image datasets.

We are randomly choosing the images to perform operations like rotation, cropping, flipping and others.

PyTorch transforms allows us to do these operations directly inside the Dataset defined by us.

Pre-processing using Transforms

random_transforms = transforms.Compose([
    transforms.RandomRotation(30),  # Randomly rotate the image by up to 30 degrees
    # transforms.RandomResizedCrop(IMAGE_SIZE),  # Randomly crop and resize the image to 224x224
    # transforms.RandomHorizontalFlip(),  # Randomly flip the image horizontally
])

# Define the fixed transformations
fixed_transforms = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5], std=[0.5])
])

# Define the overall transformation pipeline
transform = transforms.Compose([
    transforms.RandomApply([random_transforms], p=0.5),  # Apply random transformations with a probability of 0.5
    fixed_transforms
])


transform_test = transforms.Compose([
    transforms.ToTensor(),
])

Custom Dataset and Data-loaders

class SignDataSet(Dataset):
  def __init__(
      self,
      image_df, 
      label_df,
      transform,
      split = None,
  ):
    self.image_df = image_df 
    self.label_df = torch.nn.functional.one_hot(torch.tensor(np.array(label_df))).float()
    self.split = split 
    self.transform = transform

  def __len__(self):
    return len(self.label_df)
  
  def __getitem__(self, index):
    image = self.image_df.iloc[index]
    image = np.reshape(np.array(image), (28,28))

    image = Image.fromarray(image.astype(np.uint8))

    label = self.label_df[index]
    # label = torch.nn.functional.one_hot(torch.tensor(label))

    if self.split == 'train':
      image = self.transform(image)

    if self.split == 'test':
      image = self.transform(image)
    return image, label

def make_loader(x, y, config, mode):
  data = SignDataSet(x, y, transform, mode)
  data_loader = DataLoader(data, batch_size = config['batch_size'], drop_last = True)
  return data_loader

Defining and training the model

Now we can build our model. This particular model consists of two CNN layers with one output layer. This model is just for the demo purposes but we can tweak any parameters in the config defined above.

Model takes the input as 28×28 input image size and then classifies them into 25 categories.

class SignLabelModel(nn.Module):
    def __init__(self, kernels, num_classes):
        super(SignLabelModel, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Conv2d(1, kernels[0], kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2))
        
        self.layer2 = nn.Sequential(
            nn.Conv2d(kernels[0], kernels[1], kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        self.classifier = nn.Sequential(
            nn.Linear(32 * 7 * 7, 128),
            nn.ReLU(inplace=True),
            nn.Linear(128, num_classes)
        )

    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x

Setting up model, optimiser and loss function

Once we have everything read we can define the entire model in the following way.

We have set nn.CrossEntropyLoss() which will calculate the loss for the categorical data. Adam optimizer is an obvious choice for the CNN architecture.

You can even try to tweak these parameters through config.

def model_creation(config):
    train_loader = make_loader(X_train, y_train, config, 'train')
    test_loader = make_loader(X_val, y_val, config, 'test')

    # Make the model
    model = SignLabelModel(config.kernels, config.classes).to(device)

    # Make the loss and optimizer
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(
        model.parameters(), lr=config.learning_rate)
    
    return model, train_loader, test_loader, criterion, optimizer

Model training

wandb.watch

Following is a training loop for our CNN model. A different thing to note here is the inclusion of wandb.watch which tells wandb to track all the experiments which we are going to perform; log is set to all.

Entire training loop is similar to any other deep learning loop.

def train_model(model, loader, criterion, optimizer, config):
    # Tell wandb to watch what the model gets up to: gradients, weights, and more!
    wandb.watch(model, criterion, log="all", log_freq=10)

    # Run training and track with wandb
    total_batches = len(loader) * config.epochs
    example_ct = 0  # number of examples seen
    batch_ct = 0
    for epoch in tqdm(range(config.epochs)):
        for _, (images, labels) in enumerate(loader):

            loss = train_batch(images, labels, model, optimizer, criterion)
            example_ct +=  len(images)
            batch_ct += 1

            # Report metrics every 25th batch
            if ((batch_ct + 1) % 25) == 0:
                train_log(loss, example_ct, epoch)


def train_batch(images, labels, model, optimizer, criterion):
    images, labels = images.to(device), labels.to(device)
    
    # Forward pass ➡
    outputs = model(images)
    loss = criterion(outputs, labels)
    
    # Backward pass ⬅
    optimizer.zero_grad()
    loss.backward()

    # Step with optimizer
    optimizer.step()

    return loss

See the dashboard during the training with following details,

wandb.save

We can make use of wandb.save to save our model in the ONNX format. Which will then allow us to deploy it on wide set of platforms.

We even can visualise this model into the dashboard.

You would need to install onxx before this step.

! pip install onnx

Now we can test the model and save it,

def test_model(model, test_loader):
    model.eval()

    # Run the model on some test examples
    with torch.no_grad():
        correct, total = 0, 0
        for images, labels in test_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            labels = [torch.argmax(l).item() for l in labels]
            # print(predicted.shape, labels)
            # correct += (predicted == labels).sum().item()
            count = sum([1 for i in range(len(labels)) if labels[i] == predicted[i].item()])

        print(f"Accuracy of the model on the {total} " +
              f"test images: {correct / total:%}")
        
        wandb.log({"test_accuracy": correct / total})

    # Save the model in the exchangeable ONNX format
    torch.onnx.export(model, images, "model.onnx")
    wandb.save("model.onnx")
     

You can even visualise the model in dashboard.

wandb.log

As discussed, wandb.log will save all the logs and show them into the dashboard.

def train_log(loss, example_ct, epoch):
    # Where the magic happens
    wandb.log({"epoch": epoch, "loss": loss}, step=example_ct)
    print(f"Loss after {str(example_ct).zfill(5)} examples: {loss:.3f}")

You can see these logs in Wandb logs dashboard,

Running entire pipeline

Now it’s time to run the entire pipeline. Live training can be visualised in the dashboard.

model = model_pipeline(config)

After running the entire pipeline you can check out the final results in the dashboard. You can see the loss per epochs, gradients and check out the condition if the gradients are exploding or converging. These kind of things are important visualise while experimenting with deep learning models. There’s a ton more we can do with WandB but that’s for another blog!

Thanks for reading. Visit my GitHub for more projects!