Looka Engineering
Published in

Looka Engineering

How to use Tensorboard with PyTorch in Google Colab

Why you’re here

PyTorch is the fastest growing deep learning framework. It offers several benefits over the more established TensorFlow.

However, one area PyTorch falls short of TensorFlow is ecosystem support. Tensorflow has a rich ecosystem of libraries that PyTorch doesn’t have. For example, to serve models, deploy on mobile, and to visualize training. This last one is what interests me today. In particular, PyTorch doesn’t have a native training visualization tool like TensorFlow’s TensorBoard. This means it can be more time consuming to set up a visualization of your training with PyTorch than with TensorFlow, and you might decide to not set up visualization at all.

Cool Tensorflow visualisation from: https://stackoverflow.com/a/41370610/1514728

In this post I’ll show you two ways you can visualize your PyTorch model training when using Google Colab. The first uses the new Jupyter TensorBoard magic command, and the second uses the library tensorboardcolab. You can find a link to an example Colab notebook at the end of each section.

Tensorboard Colab magic

Magic commands come from the IPython kernel and are meant to concisely solve common problems in data processing. Conveniently, there’s now a TensorBoard magic command (just make sure you install the latest TensorFlow build).

# Install latest Tensorflow build
!pip install -q tf-nightly-2.0-preview
from tensorflow import summary
%load_ext tensorboard.notebook

Then instantiate the summary writers. In this case I have one for recording training and another for recording testing.

current_time = str(datetime.datetime.now().timestamp())
train_log_dir = 'logs/tensorboard/train/' + current_time
test_log_dir = 'logs/tensorboard/test/' + current_time
train_summary_writer = summary.create_file_writer(train_log_dir)
test_summary_writer = summary.create_file_writer(test_log_dir)

And you can write to TensorBoard just like you would in your TensorFlow code. See lines 24–26 below.

Then run

%tensorboard --logdir logs/tensorboard

It can take 5 seconds for TensorBoard to load and start listening to the folder logs/tensorboard. Once you see TensorBoard up, run.

train(model, train_loader, ...)

And you’ll see something like this:

Tensorboard running simultaneously with training!

Congratulations 🎉. You’re using PyTorch with TensorBoard in Colab. Note that this should also work in any Jupyter notebook using the IPython kernel as long as you install the latest TensorFlow.

To see how everything works together checkout this example Colab notebook.

The tensorboardcolab Library

The second way to use TensorBoard with PyTorch in Colab is the tensorboardcolab library. This library works independently of the TensorBoard magic command described above.

This approach is similar to the TensorBoard magic command, except instead of running TensorBoard in your colab notebook it uses ngrok to tunnel TensorBoard to localhost. See this Stack Overflow Answer for more details.

To use tensorboardcolab we’ll instantiate TensorBoardColab and then save values to it during training.

!pip install tensorboardcolabfrom tensorboardcolab import TensorBoardColab
tb = TensorBoardColab()

The code to save to TensorBoardColab while training could look something like this (see line 25):

Then when you run

train(model, train_loader, ...)

You’ll see something like

Tensorboard running on http://db797eee.ngrok.io

When you go to the link, you’ll see the TensorBoard dashboard you know and love.

Tensorboard dashboard using tensorboardcolab

To see how everything works together checkout this example Colab notebook.


There you have it, two ways to visualize your PyTorch training using TensorBoard and Google Colab. I’ve made example Colab notebooks for each of the approaches.

Magic Tensorboard: https://colab.research.google.com/drive/1NbEqqB42VSzYt-mmb4ESc8yxL05U2TIV

tensorboardcolab: https://colab.research.google.com/drive/1hR-DQvve8uEX2zH8h4y1XgP1atKRUl0g

Good luck 😄

You can follow me on Twitter for more content @andrewbrownmart




Creative engineers, data scientists, and designers making great design accessible and delightful for everyone.

Recommended from Medium

Object Detection State of the Art 2022

Monte Carlo Markov Chain (MCMC) explained

Natural Language Processing-Tokenization

# Image Detection Using Convolutional Neural Networks

Eyes Blink Detector and Counter MediaPipe

What is Machine Learning and Why It Matters in 2018

Improving Vanilla Gradient Descent

Data Product & Economy 2.0

Get the Medium app

A button that says 'Download on the App Store', and if clicked it will lead you to the iOS App store
A button that says 'Get it on, Google Play', and if clicked it will lead you to the Google Play store
Andrew B. Martin

Andrew B. Martin

CTO @ Looka | @andrewbrownmart

More from Medium

Adventures in Deploying a Deep Learning Model in the Browser

I Converted the normal Image into cartoonish with OpenCV in Python

Selecting closest audio to given audio using machine learning algorithms

Solving GANs: NN architecture & feedback