Make TensorBoard Work with PyTorch

Source: Deep Learning on Medium

TL;DR

Say goodbye to your nightly version of PyTorch, and update it to version 1.4.

Problem: graphs are not showing in TensorBoard

When you use the TensorBoard with versions of PyTorch, like 1.1, 1.2, etc., you might encounter a problem, the graphs are not showing in TensorBoard.

This problem is discussed in the issues. But the solutions are outdated, like downloading the PyTorch nightly version, using conda install, or using the tensorboardX. After doing some research, I found this problem is solved by the latest version of PyTorch 1.4. Below is a working example.

Make TensorBoard Work with PyTorch

1 Install libraries

# python 3.6 environment
pip install torch==1.4.0 torchvision==0.5.0 tensorboard==2.1.0

2 Write demo code

Saving the code to tensorboard_demo.py

tensorboard_demo.py

3 TensorBoard demo

Open the terminal and execute the below commands.

python tensorboard_demo.py
tensorboard --logdir=runs

Then opening the http://localhost:6006/ in the browser, the graph should be shown.

Check out my other posts on Medium, with a better TOC list view!
GitHub: https://github.com/BrambleXu
LinkedIn:
www.linkedin.com/in/xu-liang
Blog:
https://bramblexu.org

Reference