How to run Tensorboard for PyTorch 1.1.0 inside Jupyter notebook

(Comments)

tb_pytorch_nb

Facebook introduced PyTorch 1.1 with TensorBoard support. Let's try it out really quickly on Colab's Jupyter Notebook.

Not need to install anything locally on your development machine. Google's Colab cames in handy free of charge even with its upgraded Tesla T4 GPU.

Firstly, let's create a Colab notebook or open this one I made.

Type in the first cell to check the version of PyTorch is at minimal 1.1.0

import torch
torch.__version__

Then you are going to install the cutting edge TensorBoard build like this.

!pip install -q tb-nightly

The output might remind you to restart the runtime to make the new TensorBoard take effect. You can click through Runtime -> Restart runtime....

Next, load the TensorBoard notebook extension with this magic line.

%load_ext tensorboard

After which you can start by exploring the TORCH.UTILS.TENSORBOARD API, these utilities let you log PyTorch models and metrics into a directory for visualization within the TensorBoard UI. Scalars, images, histograms, graphs, and embedding visualizations are all supported for PyTorch models and tensors.

The SummaryWriter class is your main entry to log data for consumption and visualization by TensorBoard. Let's run this official demo for MNIST dataset and ResNet50 model.

import torch
import torchvision
from torch.utils.tensorboard import SummaryWriter
from torchvision import datasets, transforms

# Writer will output to ./runs/ directory by default
writer = SummaryWriter()

transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
trainset = datasets.MNIST('mnist_train', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)
model = torchvision.models.resnet50(False)
# Have ResNet model take in grayscale rather than RGB
model.conv1 = torch.nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
images, labels = next(iter(trainloader))

grid = torchvision.utils.make_grid(images)
writer.add_image('images', grid, 0)
writer.add_graph(model, images)
writer.close()

You just wrote an image and the model graph data to TensorBoard summary. The writer wrote the output file to "./runs" directory by default.

Let's run the TensorBoard to visualize them

%tensorboard --logdir=runs

That's it, you have it!

tb_pytorch

Summary and Further reading

This really short tutorial gets you to start with running TensorBoard with latest Pytorch 1.1.0 in a Jupyter Notebook. Keep playing around with other features supported with PyTorch TensorBoard.

Read the official API document here - TORCH.UTILS.TENSORBOARD

Current rating: 3.9

Comments