PyTorch

PyTorch is a deep learning framework that is extremely popular among researchers. We have written extensively about the differences between TensorFlow and PyTorch.

In this section, we will look at how TensorBoard can be used from PyTorch. Let's use a specific example to describe the steps. In this example, we will use the popular MNIST dataset. With TensorFlow, we used the Summary API to create the file writer that would log data into the logdir folder. With PyTorch, there is a similar file writer. Here is a code snippet showing how to use it.

# Import the summary writer
from torch.utils.tensorboard import SummaryWriter
# Create an instance of the object
writer = SummaryWriter()

Once the file writer is created, we will write the data (i.e. image etc) to TensorBoard. See the section highlighted in the code snippet below.

import torch
import torchvision
from torchvision import datasets, transforms
# Compose a set of transforms to use later on
transform = transforms.Compose([
    transforms.ToTensor(), 
    transforms.Normalize((0.5,), (0.5,))
])
# Load in the MNIST dataset
trainset = datasets.MNIST(
    'mnist_train', 
    train=True, 
    download=True, 
    transform=transform
)
# Create a data loader
trainloader = torch.utils.data.DataLoader(
    trainset, 
    batch_size=64, 
    shuffle=True
)
# Get a pre-trained ResNet18 model
model = torchvision.models.resnet18(False)
# Change the first layer to accept grayscale images
model.conv1 = torch.nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
# Get the first batch from the data loader
images, labels = next(iter(trainloader))

# Write the data to TensorBoard
grid = torchvision.utils.make_grid(images)
writer.add_image('images', grid, 0)
writer.add_graph(model, images)
writer.close()

Once the user navigates to TensorBoard, they will get output similar to what they would see with TensorFlow i.e. users can use the summary writer to write data just like they would with the Summary API.

---