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.
After running this code, launch TensorBoard against the log directory and open the Images tab.
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.
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.
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:
- '
uint8with values from0to255' - 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.
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
Nsteps 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.imageto 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
stepvalue makes it difficult to compare image evolution during training.
Summary
- Use
tf.summary.imageto 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
stepvalue. - 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.

