3 simple steps to use PyTorch & TensorBoard on Google Colab in 2024

I recently attempted to visualize a PyTorch training session on a Google Colab notebook using TensorBoard. Since all the tutorials and samples I found were slightly outdated, I thought it would be a good idea to share a solution that works as of March 2024.

Currently, you no longer need to install additional libraries nor use ngrok to tunnel the TensorBoard port as you did in the past. Additionally, Colab now comes with TensorFlow 2.x by default.

Step 1: Load the TensorBoard extension

The TensorBoard notebook extension allows you to use TensorBoard inside your notebook without having to manage runtime, ports, etc.

%load_ext tensorboard

Step 2: Run TensorBoard

Finally you need to initialize and start TensorBoard within the notebook using the magic command %tensorboard. It takes the same arguments as the command line invocation.

%tensorboard --logdir runs

Step 3: PyTorch’s TensorBoard tutorial

Now, you are basically good to go. You can follow PyTorch’s TensorBoard tutorial which teaches you all the basic stuff.

Thanks for reading!