TensorBoard: Visualizing Training Process in TensorFlow 2.0

Written by / Li Xihan, Google Developers Expert

This article is excerpted from “Simple and Rough TensorFlow 2.0”

First, create a folder in the code directory (e.g., ./tensorboard) to store the TensorBoard log files, and instantiate a logger in the code:

summary_writer = tf.summary.create_file_writer('./tensorboard')     # The parameter is the directory where the log files are saved

summary_writer = tf.summary.create_file_writer('./tensorboard')
# Start model training
for batch_index in range(num_batches):
    # ... (training code, current batch loss value stored in variable loss)
    with summary_writer.as_default():                               # The logger you want to use
        tf.summary.scalar("loss", loss, step=batch_index)
        tf.summary.scalar("MyScalar", my_scalar, step=batch_index)  # You can also add other custom variables

Every time tf.summary.scalar() is called, the logger writes a record to the log file. In addition to the simplest scalars, TensorBoard can also visualize other types of data (such as images, audio, etc.), see the TensorBoard documentation for details.

When we want to visualize the training process, open a terminal in the code directory (enter the TensorFlow conda environment if necessary) and run:

tensorboard --logdir=./tensorboard

Then use a browser to access the URL output by the command line program (usually http://computer_name:6006) to access the TensorBoard visual interface, as shown in the figure below:

By default, TensorBoard updates data every 30 seconds. However, you can also manually refresh by clicking the refresh button in the upper right corner.

There are a few points to note when using TensorBoard:

• If you need to retrain, you need to delete the information in the log folder and restart TensorBoard (or create a new log folder and start TensorBoard, setting the --logdir parameter to the newly created folder);

• Keep the log folder directory entirely in English.

Finally, here is an example using the multilayer perceptron model from the previous chapter to demonstrate the use of TensorBoard:

import tensorflow as tf
from zh.model.mnist.mlp import MLP
from zh.model.utils import MNISTLoader

num_batches = 10000
batch_size = 50
learning_rate = 0.001
model = MLP()
data_loader = MNISTLoader()
optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate)
summary_writer = tf.summary.create_file_writer('./tensorboard')     # Instantiate logger
for batch_index in range(num_batches):
    X, y = data_loader.get_batch(batch_size)
    with tf.GradientTape() as tape:
        y_pred = model(X)
        loss = tf.keras.losses.sparse_categorical_crossentropy(y_true=y, y_pred=y_pred)
        loss = tf.reduce_mean(loss)
        print("batch %d: loss %f" % (batch_index, loss.numpy()))
        with summary_writer.as_default():                           # Specify logger
            tf.summary.scalar("loss", loss, step=batch_index)       # Write the current loss value to the logger
    grads = tape.gradient(loss, model.variables)
    optimizer.apply_gradients(grads_and_vars=zip(grads, model.variables))