How to use TensorBoard with PyTorch

TensorBoard (Image Source: TensorFlow)

TensorBoard is a tool for visualizing and understanding the performance of deep learning models. It is an open-source tool developed by TensorFlow and can be used with any deep learning framework. TensorBoard allows tracking and visualizing metrics such as loss and accuracy, visualizing the model graph, viewing histograms, displaying images and optimize your TensorFlow/PyTorch codes.

PyTorch Installation

For the usage of TensorBoard with PyTorch, the installation of PyTorch should be installed to log models and metrics into TensorBoard log directory.

The following command will install PyTorch 1.4+ via Anaconda (recommended):

$ conda install pytorch torchvision -c pytorch

or pip

$ pip install torch torchvision

TensorBoard Installation

Install TensorBoard through the command line to visualize data you logged.

$ pip install tensorboard

Use TensorBoard with PyTorch

In order to use TensorBoard with PyTorch, you need to create a SummaryWriter instance to ensure all logs are created and stored.

The SummaryWriter class provides a high-level API to create an event file in a given directory and add summaries and events to it. The class updates the file contents asynchronously. This allows a training program to call methods to add data to the file directly from the training loop, without slowing down training.

import torch
from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter()

Writer will output to ./runs/ directory by default.

If you do not need the summary writer anymore, call close() method.

writer.close()

Steps to add TensorBoard

Adding TensorBoard to your PyTorch model will take a few simple steps:

1. Starting with a simple Neural Network.
2. Initializing the SummaryWriter which allows us to write to TensorBoard.
3. Writing away some scalar values, both individually and in groups.
4. Writing away images, graphs and histograms.

Run TensorBoard

Start TensorBoard by specifying the root log directory you used above. Argument logdirpoints to directory where TensorBoard will look to find event files that it can display. TensorBoard will recursively walk the directory structure rooted at logdir, looking for .*tfevents.* files.

Load the TensorBoard notebook extension:

%load_ext tensorboard

$ tensorboard --logdir=runs\

Go to the URL it provides OR to http://localhost:6006/

Example of use

To show the use of TensorBoard we will run the Convolutional Neural Network model on the Fashion-MNIST dataset. torchvision already has the Fashion-MNIST dataset.

If you’re not familiar with Fashion-MNIST dataset:

Fashion-MNIST is a dataset of Zalando's article images—consisting of a training set of 60,000 examples and a test set of 10,000 examples. Each example is a 28x28 grayscale image, associated with a label from 10 classes. We intend Fashion-MNIST to serve as a direct drop-in replacement for the original MNIST dataset for benchmarking machine learning algorithms. It shares the same image size and structure of training and testing splits. — From Github

Fashion-MNIST Dataset — From GitHub

Network

import torch 
import torchvision

import numpy as np 

EPOCHS = 10 
BATCH_SIZE = 64

xy_trainPT = torchvision.datasets.MNIST(root='./data', train=True, download=True,
                                        transform=torchvision.transforms.Compose([torchvision.transforms.ToTensor()])) 
xy_trainPT_loader = torch.utils.data.DataLoader(xy_trainPT, batch_size=BATCH_SIZE)


def model(hidden):
    model= torch.nn.Sequential(            
         torch.nn.Linear(784,hidden),
         torch.nn.Sigmoid(), 
         torch.nn.Linear(hidden,10), 
         torch.nn.LogSoftmax(dim=1) 
         )
    return model

Import SummaryWriter class:

from torch.utils.tensorboard import SummaryWriter%load_ext tensorboard

Create SummaryWriter instance and train the model:

modelPT = model(10)
criterion = torch.nn.NLLLoss()
optimizer = torch.optim.SGD(modelPT.parameters(), lr=0.01)

writer = SummaryWriter()
for e in range(EPOCHS):
    running_loss = 0
    for images, labels in xy_trainPT_loader:
        images = images.view(images.shape[0], -1)
        output = modelPT(images)
        loss = criterion(output, labels)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
        running_loss += loss.item()
    print("Epoch {} - Training loss: {}".format(e, running_loss/len(xy_trainPT_loader)))
    writer.add_scalar("loss x epoch", running_loss/len(xy_trainPT_loader), e)

 writer.close()

Load the TensorBoard notebook extension and run TensorBoard

%load_ext tensorboard

$ tensorboard --logdir=runs\

TensorBoard shows how the loss and accuracy change with every epoch

TensorBoard Dashboard

The TensorBoard dashboard shows how the loss and accuracy change with every epoch. You can use it to also track training speed, learning rate, and other scalar values. It’s helpful to compare these metrics across different training runs to improve your model.

Share TensorBoard dashboards

TensorBoard.dev lets you upload and share your ML experiment results with anyone. Use TensorBoard.dev to host, track, and share your TensorBoard dashboards.

Install the latest version of TensorBoard to use the uploader.

$ pip install tensorboard --upgrade

Use a simple command to upload and share your TensorBoard.

$ tensorboard dev upload --logdir runs \
--name "My latest experiment" \ # optional
--description "Simple comparison of several hyperparameters" # optional

TensorBoard.dev currently supports scalars, graphs, histograms, distributions, hparams, and text dashboards.

Note: Uploaded TensorBoards are public and visible to everyone. Do not upload sensitive data.

Conclusion

TensorBoard is a powerful tool for visualizing and understanding the performance of deep learning models. With TensorBoard, you can gain insights into the behavior of your model that would be difficult to obtain by simply looking at the code. Whether you are a beginner or an experienced deep learning practitioner, TensorBoard is a tool that you should consider using to help you better understand your models.

Stay tune to Part 2 where we will explore how to use TensorBoard to optimize hyperparameters for model training.