Tensorboard
Custom Images
Matplotlib
Machine Learning
Visualization

How to Display Custom Images in Tensorboard e.g. Matplotlib Plots?

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Introduction

TensorBoard is not limited to scalar metrics such as loss and accuracy. It can also display images, which makes it useful for visualizing model inputs, generated outputs, segmentation masks, or custom Matplotlib plots during training.

Core Sections

Log image tensors directly with tf.summary.image

The core API for image logging is tf.summary.image. It expects a batch of images, usually shaped as batch x height x width x channels.

python
1import tensorflow as tf
2
3writer = tf.summary.create_file_writer("logs/images")
4
5image = tf.random.uniform(shape=(1, 64, 64, 3), maxval=255, dtype=tf.int32)
6image = tf.cast(image, tf.uint8)
7
8with writer.as_default():
9    tf.summary.image("sample_image", image, step=0)

After running this code, launch TensorBoard against the log directory and open the Images tab.

bash
tensorboard --logdir logs/images

This is the most direct option when the data already exists as image tensors.

Convert a Matplotlib figure into an image

If you want to display a plot rather than a raw image tensor, render the Matplotlib figure into an in-memory PNG and decode it back into a TensorFlow image.

python
1import io
2import matplotlib.pyplot as plt
3import tensorflow as tf
4
5
6def plot_to_image(figure):
7    buffer = io.BytesIO()
8    figure.savefig(buffer, format="png")
9    plt.close(figure)
10    buffer.seek(0)
11
12    image = tf.image.decode_png(buffer.getvalue(), channels=4)
13    image = tf.expand_dims(image, 0)
14    return image
15
16
17writer = tf.summary.create_file_writer("logs/plots")
18
19figure, ax = plt.subplots()
20ax.plot([1, 2, 3, 4], [1, 4, 9, 16])
21ax.set_title("Training Example Plot")
22
23plot_image = plot_to_image(figure)
24
25with writer.as_default():
26    tf.summary.image("matplotlib_plot", plot_image, step=0)

This pattern is common when you want to inspect confusion matrices, prediction overlays, or custom diagnostics that are easier to produce with Matplotlib than with raw tensors.

Log images during training

The usual place to write custom images is inside a training loop or callback. For example, you might log a generated sample every few epochs.

python
1for step in range(5):
2    image = tf.random.uniform(shape=(1, 32, 32, 1))
3
4    with writer.as_default():
5        tf.summary.image("generated_sample", image, step=step)

The step value is important because TensorBoard uses it to place images on the training timeline. Without a meaningful step, the logged outputs are much harder to compare over time.

Keep image shape and dtype under control

TensorBoard expects actual image-like tensors. Common safe choices are:

  • 'uint8 with values from 0 to 255'
  • floating-point tensors normalized to a valid image range

If your data is not already shaped as a batch, add the batch dimension with tf.expand_dims. If the image has one channel, keep the last dimension so the tensor is still rank 4.

python
single_image = tf.random.uniform((64, 64, 1))
batched = tf.expand_dims(single_image, axis=0)

Most logging issues come from incorrect rank or unexpected dtype rather than from TensorBoard itself.

Be selective about logging frequency

Image summaries are much heavier than scalar summaries. Logging too many full-resolution images too often will create large event files and slow down both training and TensorBoard.

A practical strategy is:

  • log only a few representative images
  • log every N steps or epochs rather than every batch
  • resize very large images before writing them

This keeps the visualization useful without turning the log directory into a storage problem.

Common Pitfalls

  • Forgetting the batch dimension causes tf.summary.image to receive the wrong tensor rank.
  • Logging Matplotlib figures without converting them to image tensors first does not work because TensorBoard expects image data, not figure objects.
  • Writing too many large images too frequently produces oversized event files and slow dashboard performance.
  • Using unexpected dtypes or value ranges can make images appear blank, clipped, or incorrectly colored.
  • Omitting a meaningful step value makes it difficult to compare image evolution during training.

Summary

  • Use tf.summary.image to display custom images in TensorBoard.
  • If the source is a Matplotlib plot, render it to PNG in memory and decode it as an image tensor.
  • Always provide the batch dimension and a sensible step value.
  • Watch dtype, shape, and value range to avoid malformed image summaries.
  • Log images selectively so the visualization remains useful and the event files stay manageable.

Course illustration
Course illustration

All Rights Reserved.