Machine learning is complicated. There are countless options available and a lot to track. Fortunately, there’s TensorBoard, which makes the process easy.
When developing machine learning models, there are many factors: How many epochs for training, the loss metric, or even the model structure. Each of these decisions can propagate and make the difference between an ineffective model and a successful model.
This article will discuss some features in TensorBoard and how you can set up TensorBoard for your next machine learning project. In particular, the focus will be on using TensorBoard with TensorFlow and Keras-based models.
To determine an optimal configuration of your model, you need to run experiments. And as any data scientist knows, you need to track and evaluate these experiments effectively.
Fortunately, TensorBoard has a lot of built-in functionality that you can utilize to understand what is going on within your model quickly.
Setup
Setting up TensorBoard is a simple process. With only a few lines of code, you’re able to track key metrics from your machine learning model.
The first step is to load the extension within your notebook.
| # Load the TensorBoard notebook extension | |
| %load_ext tensorboard |
Load TensorBoard Extension (Code by Author)
Next, you’ll need some data. Here I’m using the MNIST dataset built into TensorFlow. The data is reshaped to allow for the use of 2D convolutional layers.
| import tensorflow as tf | |
| import datetime | |
| mnist = tf.keras.datasets.mnist | |
| (x_train, y_train),(x_test, y_test) = mnist.load_data() | |
| x_train, x_test = x_train / 255.0, x_test / 255.0 | |
| x_train = x_train.reshape(60000,28,28,1) | |
| x_test = x_test.reshape(10000,28,28,1) |
Load MNIST dataset from TensorFlow (Code by Author)
After the data is prepared, you need to build your model. Here I’ve gone a bit overboard and parameterized each variable.
This choice is made to quickly change different aspects of the model without too much issue. This structure also avoids you having values where it isn’t clear what the value is changing.
The structure is simple. But this basic structure shows a variety of different layers, which you can analyze in TensorBoard. The model contains a 2D convolutional layer followed by a pooling layer, dropout layer, output flattening, and a dense layer.
There is also an adjustment to the learning rate. In contrast to setting a fixed learning rate, the scheduler makes the learning rate flexible and helps with convergence.
| DENSE_UNITS = 20 | |
| DROPOUT = 0.3 | |
| ACTIVATION_INTERNAL = ‘relu’ | |
| ACTIVATION_FINAL = ‘softmax’ | |
| NUM_FILTERS = 32 | |
| POOL_SIZE = (2, 2) | |
| KERNEL_SIZE = (3,3) | |
| INPUT_SHAPE = (28,28,1) | |
| EPOCHS = 100 | |
| LOSS = ‘sparse_categorical_crossentropy’ | |
| METRICS = [‘accuracy’] | |
| LEARNING_RATE_INITIAL = 1e-2 | |
| LEARNING_DECAY_STEPS = 10000 | |
| LEARNING_DECAY_RATE = 0.9 | |
| LR_SCHEDULE = tf.keras.optimizers.schedules.ExponentialDecay( | |
| initial_learning_rate=LEARNING_RATE_INITIAL, | |
| decay_steps=LEARNING_DECAY_STEPS, | |
| decay_rate=LEARNING_DECAY_RATE) | |
| OPTIMIZER = tf.keras.optimizers.Adam(learning_rate=LR_SCHEDULE) | |
| def create_model(): | |
| return tf.keras.models.Sequential([ | |
| tf.keras.layers.Conv2D(NUM_FILTERS, kernel_size=KERNEL_SIZE, activation=ACTIVATION_INTERNAL, input_shape=INPUT_SHAPE, name=’Conv2D_Layer’), | |
| tf.keras.layers.MaxPooling2D(pool_size=POOL_SIZE, name=’Pooling_Layer’), | |
| tf.keras.layers.Dropout(DROPOUT, name=’Dropout_Layer’), | |
| tf.keras.layers.Flatten(name=’Flatten_Layer’), | |
| tf.keras.layers.Dense(DENSE_UNITS, activation=ACTIVATION_FINAL, name=’Dense_Layer’) | |
| ]) |
view rawmodel_creation.py hosted with ❤ by GitHub
Model Creation and Parameter Definitions (Code by Author)
To add TensorBoard functionality to your existing Keras-based TensorBoard model, you need to add a callback during the model fit phase of training.
Histogram computation should be enabled to track progress effectively, and this is done by setting the historgram_freq parameter to 1.
The callback function requires a log directory to store the results of training the model. Therefore, it is beneficial to include some structured ordering in your logs for future reference. The current time is used here.
| model = create_model() | |
| model.compile(optimizer=OPTIMIZER, loss=LOSS, metrics=METRICS) | |
| log_dir = “logs/fit/” + datetime.datetime.now().strftime(“%Y%m%d-%H%M%S”) | |
| tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=log_dir, histogram_freq=1) | |
| model.fit(x=x_train, | |
| y=y_train, | |
| epochs=EPOCHS, | |
| validation_data=(x_test, y_test), | |
| callbacks=[tensorboard_callback]) |
Compile and Fit Model, Setup TensorBoard Callback (Code by Author)
Once the model is created, compiled, and fitted, the logs should be packed. Full of all the details of your model during training. Ready for you to analyze.
To view the training process results within TensorBoard, all that’s left is to run the extension.
| %tensorboard –logdir logs/fit |
Run TensorBoard (Code by Author)
Components of TensorBoard
TensorBoard is broken down into several components. These components allow you to track different metrics such as accuracy, root mean squared error or log loss. They also allow for visualization of the model as a graph and a whole lot more.
I’m showing the scalars, graphs, distributions, histograms, and time-series tabs in this post. But a list of the other available views is found in the inactive dropdown.
TensorBoard also has some styling options. I’ve used the dark mode for some of the images here.
Scalars
Scalars are the first tab you will see when opening TensorBoard. The focus here is on the performance of the model across multiple epochs.
Both the models’ loss function and any metrics that you’ve tracked are shown here.
An essential feature of this tab is the smoothing function. When dealing with many epochs or a precarious model, the overall trend can be lost. Therefore, you want to make sure that your model is improving during training and not stagnating.
By increasing the smoothing, you can view the overall trends of the model during the training process.
The scalars tab is crucial for identifying when a model is overfitting. For example, when your training metric keeps improving but there isn’t an increase in the validation plot, you may be overfitting on the validation set.

Graphs
The graphs tab allows you to view the structure of the model you’ve created. Essentially, it shows what’s happening behind the scenes.
These details are helpful when you need to share the structure of the graphs with others. In addition, the ability to upload or download graphs is available.
In addition to the base model structure, the graph also reveals how different metrics are used and the optimizer.
Here I’ve selected the sequential node with the graph. Once selected, the models’ structure is shown. The details are visible within the red box in the image below.

Distributions & Histograms
The distributions and histograms tabs are pretty similar. However, they allow you to view the same information through different visualizations.
The distributions tab gives you a good overview of the changes in the model’s weights over time. This perspective serves as an initial gauge to see if something has gone wrong.
The histograms view gives a more detailed breakdown of the exact values learned by your model.
These two visualizations are used to determine when the model is over-relying on a small set of weights. Or if the weights converge over many epochs.
Distributions

Histograms

Time-Series
The last tab shown here in TensorBoard is the time-series tab.
This view is quite similar to the scalars view. However, one distinction is the observations of your target metric for each iteration of training instead of each epoch.
Observing the model training in this manner is much more granular. This type of analysis is best when the model is not converging, and the progress over epochs is not revealing any answers.

Wrap Up
TensorBoard is a powerful tool. Through several different components and views, you can rapidly analyze your machine learning and deep learning models.
The tool is easy to set up and provides valuable insights into how to train your model better.
Here I’ve shown you only a sample of what is possible with TensorBoard. Feel free to copy the code and explore the tool for yourself.