Using TensorBoard with PyTorch
Welcome to this neural network programming series. In this episode, we will learn how to use TensorBoard to visualize metrics of our CNN during the neural network training process.
Without further ado, let's get started.
At this point in the series, we've just finished running our network through the training process. Now, we want to get more metrics about this process to better understand just what's going on under the hood.
Bird's eye view of where we are in the training process.
- Prepare the data
- Build the model
- Train the model
Analyze the model's results
- Using TensorBoard for this
TensorBoard: TensorFlow's Visualization Toolkit
TensorBoard provides the visualization and tooling needed for machine learning experimentation:
- Tracking and visualizing metrics such as loss and accuracy
- Visualizing the model graph (ops and layers)
- Viewing histograms of weights, biases, or other tensors as they change over time
- Projecting embeddings to a lower dimensional space
- Displaying images, text, and audio data
- Profiling TensorFlow programs
- And much more
As of PyTorch version
1.1.0, PyTorch has added a
tensorboard utility package that enables us to use TensorBoard with PyTorch.
from torch.utils.tensorboard import SummaryWriter
Installing TensorBoard for PyTorch
To install TensorBoard for PyTorch, use the following steps:
Verify that you are running PyTorch version
Verify that you are running TensorBoard version
Note that the TensorBoard that PyTorch uses is the same TensorBoard that was created for TensorFlow. Check the version of TensorBoard installed on your system using the this command:
At the time of the writing of this blog, TensorBoard version
1.15is not yet released. However, we can install the nightly build using the following command.
pip install tb-nightly
After getting TensorBoard version
1.15or greater installed, we're ready to go!
Note that the PyTorch docs say that TensorBoard version
1.14 is the requirement. However, I was unable to the the full functionality to work on the
1.14 release. This is why
the nightly build is being used in the video.
Getting Started with TensorBoard for PyTorch
TensorBoard is a font-end web interface that essentially reads data from a file and displays it. To use TensorBoard our task is to get the data we want displayed saved to a file that TensorBoard can read.
To make this easy for us, PyTorch has created a utility class called
SummaryWriter. To get access to this class we use the following import:
from torch.utils.tensorboard import SummaryWriter
Once we have imported the class, we can create an instance of the class that we'll then use to get the data out of our program and onto the file system where it can then be consumed by TensorBoard.
Network Graph and Training Set images
SummaryWriter class comes with a bunch of method that we can call to selectively pick and choose which data we want to be available to TensorBoard. We'll start by first by passing
our network and a batch of images to the writer.
tb = SummaryWriter() network = Network() images, labels = next(iter(train_loader)) grid = torchvision.utils.make_grid(images) tb.add_image('images', grid) tb.add_graph(network, images) tb.close()
This code creates a
SummaryWriter instance called
tb for TensorBoard. Then, creates an instance of our PyTorch network and unpacks a batch of images and labels from our
PyTorch data loader object.
Then, the images and the network are added to the file that TensorBoard will consume. Effectively, we can say that the network graph and the batch of images have both been added to TensorBoard.
To launch TensorBoard, we need to run the tensorboard command at our terminal. This will launch a local server that will serve the TensorBoard UI and the the data our
By default, the PyTorch SummaryWriter object writes the data to disk in a directory called
./runs that is created in the current working directory.
When we run the tensorboard command, we pass an argument that tells tensorboard where the data is. So it's like this:
The TensorBoard server will launch and be listening for
http requests on port
6006. These details will be displayed in the console.
Access the TensorBoard UI by browsing to:
Here, we will be able to see our network graph and our image data. At the current moment, this does provide us with a visual, but it's not as useful as what comes next.
TensorBoard Histograms and Scalars
The next import types of data we can add to TensorBoard is numerical data. We can add scalar values that will be displayed over time or over epoch. We can also add values to histograms to see frequency distributions of values.
To add scalars and histograms we use the corresponding methods provided by the PyTorch
Here is an example of the calls:
tb.add_scalar('Loss', total_loss, epoch) tb.add_scalar('Number Correct', total_correct, epoch) tb.add_scalar('Accuracy', total_correct / len(train_set), epoch) tb.add_histogram('conv1.bias', network.conv1.bias, epoch) tb.add_histogram('conv1.weight', network.conv1.weight, epoch) tb.add_histogram('conv1.weight.grad', network.conv1.weight.grad, epoch)
And here is an example of where we would place these calls inside our training loop:
network = Network() train_loader = torch.utils.data.DataLoader(train_set, batch_size=100) optimizer = optim.Adam(network.parameters(), lr=0.01) images, labels = next(iter(train_loader)) grid = torchvision.utils.make_grid(images) tb = SummaryWriter() tb.add_image('images', grid) tb.add_graph(network, images) for epoch in range(1): total_loss = 0 total_correct = 0 for batch in train_loader: # Get Batch # Pass Batch # Calculate Loss # Calculate Gradient # Update Weights tb.add_scalar('Loss', total_loss, epoch) tb.add_scalar('Number Correct', total_correct, epoch) tb.add_scalar('Accuracy', total_correct / len(train_set), epoch) tb.add_histogram('conv1.bias', network.conv1.bias, epoch) tb.add_histogram('conv1.weight', network.conv1.weight, epoch) tb.add_histogram( 'conv1.weight.grad' ,network.conv1.weight.grad ,epoch ) print( "epoch", epoch, "total_correct:", total_correct, "loss:", total_loss ) tb.close()
This will add these values to TensorBoard. The values even update in real-time as the network trains.
It is helpful to see the loss and accuracy values over time. However, we might need to admit that TensorBoard really isn't needed for this.
The real power of TensorBoard is its out-of-the-box capability of comparing multiple runs. This allows us to rapidly experiment by varying the hyperparameter values and comparing runs to see which parameters are working best.
Hyperparameter Experimentation is Next
We should now have a good understanding of what TensorBoard is and how we can use it. In the next episode, we'll see how to leverage TensorBoard to better evaluate different network training runs. See you in the next one!