How to Use TensorBoard in TF2

Visualise and boost your trainings !

Published in

The Startup

6 min readNov 9, 2020

Over the recent years the main deep learning libraries managed to become more and more easy to use. It is now possible to create a deep learning model in just a few lines of code if not one in some cases. Training these algorithms on the other hand can still remain tricky.

“An image is worth a thousand words.”

This is even more true when we talk about machine learning (ML). Visualisation is a great interpretation of what is happening during the learning phase. Who doesn’t wish to be able to monitor the training losses, the weights and many other useful informations in order to improve the model performances. All of this can be done using TensorBoard.

TensorBoard is part of TensorFlow but it can be used with other libraries such as PyTorch. It’s a visualisation toolkit which comes with various functionalities to display different aspects of an ML pipeline.

In this article, we will cover the basic of TensorBoard and the creation of custom callbacks to track what’s best for you and get the most out of your trainings. Below are the actual sections we will cover :

How to start TensorBoard ?
How to display your training logs into TensorBoard ?
Going Further : Log what’s best for you !
Training a basic model on the Fashion MNIST classification task

How to start your TensorBoard ?

First you need to install it in your python environment. Note that if you already have TensorFlow installed it should be done, but you could check :

pip show tensorboard

If it’s not installed simply do one of these :

# option 1 : using pip
pip install tensorboard# option 2 : using conda 
conda install -n <env_name> tensorboard

Now run the command below in your terminal :

tensorboard --logdir <log_directory> --host <0.0.0.0> --port <6006>

The log_directory is the root directory where you store all your training logs but we will see this in more details with the coming sections. The default host is usually 0.0.0.0 which corresponds to your localhost and the default port is 6006.

Then open the browser of your choice and go to : 127.0.0.1:6006

Here is your TensorBoard ! Amazing isn’t it ? If it’s your first time it should be empty because you didn’t store any informations in your log_directory yet.

Just for your information, you now have the option to run TensorBoard directly in a notebook by using the magic command :

%tensorboard --logdir <log_directory>

How to display your training logs into TensorBoard ?

For now we will keep it simple and cover the basics. Let’s say that we only want to monitor the training/validation evolution for the loss and the accuracy.

Option 1. Training with the Keras API :

In this case, we just need to create a TensorBoard callback which ensures that the informations are logged in the specified directory. This directory must be inside our root log_directory mentioned earlier: log_directory/our_run/. Why ? This way we can visualise all our runs.

Go back to your board and refresh, then you should see the loss curve and the accuracy curve of both training and validation. You can also go to the graph tab and see the computational graph of your algorithm.

That wasn’t difficult was it ? We simply used the prebuilt TensorBoard callback, here are all the available callbacks you can use straight away.

Option 2. Training with a custom loop :

If you define your training loop manually, then you also need to log information manually. Hopefully the tf.summary API makes it easy.

You just need to create file writers which points to the log directories (one for train and one for validation), then with this file writer instances, you can log the desired information after every epoch inside the training loop.

You should have the exact same results as the one obtained using the Keras API.

Going further : Log what’s best for you !

The summaries doesn’t have to be scalars, you can log images, weights and bias histograms, computational graphs, computing profile, embeddings projection, texts or even audio files. It also contains specific visualisation tools to do hyper parameters selection and some more useful plugins.

To make it short you have a lot of possibilities and it’s up to you to choose what best fits your case.

Let’s take an example of a training using TensorBoard together !

Training a basic model on the Fashion MNIST classification task :

First let’s select the final Hyper Parameters

For the purpose of this article let’s say we still hesitate between some hyper parameters such as the number of units, the dropout rate and the optimizer of our classifier. We could do an experiment to see which combination is better :

Once you ran the code above, you should see eight new runs of one epoch each in your scalar tab. It would be possible to compare the performances this way but if you check the tabs you will see a new one called Hparams and it should look like this :

The parallel coordinates view clearly indicated that we should use a model with 32 units, .10 dropout rate and the Adam optimizer.

Train and log in TensorBoard :

Now that we have our hyper parameters selected let’s launch a real training and see how it goes. Here we are doing a classification task so it would be nice to log the confusion matrix after each epoch.

First, like we did before, let’s create our basic callback to monitor the loss and the accuracy evolution :

Then we need to create a callback to log our confusion matrix after each epoch :

And let’s get crazy and add a learning rate scheduler just to see how simple it is :

Finally we can merge all our callbacks and train the model, the full for the classification task and the monitoring is below. For the purpose of having all information in TB let’s log some data samples too :

Now you can go back to your TensorBoard to see the metrics, the learning rate curve and the confusion matrices (the images can take a little time to load in TB):