Introduction

TensorBoard can display more than scalar metrics. Logging custom images during training helps you inspect data quality, augmentation behavior, and model outputs in a visual workflow. In Keras, the clean approach is to write image summaries from a callback.

Set Up TensorBoard Logging in Keras

Start with a log directory and the standard TensorBoard callback for scalars. Then add a custom callback for images.

1import os
2import tensorflow as tf
3
4log_dir = "logs/image_demo"
5os.makedirs(log_dir, exist_ok=True)
6
7tb_callback = tf.keras.callbacks.TensorBoard(
8    log_dir=log_dir,
9    histogram_freq=0,
10    write_graph=True,
11)

Image summaries require tensors shaped as batch, height, width, channels. Values should be in a valid display range, usually zero to one for float inputs.

Custom Callback for Image Summaries

The callback below logs a fixed validation batch each epoch and writes model predictions as images.

1import tensorflow as tf
2
3class ImageSummaryCallback(tf.keras.callbacks.Callback):
4    def __init__(self, sample_images, writer):
5        super().__init__()
6        self.sample_images = sample_images
7        self.writer = writer
8
9    def on_epoch_end(self, epoch, logs=None):
10        preds = self.model.predict(self.sample_images, verbose=0)
11        preds = tf.clip_by_value(preds, 0.0, 1.0)
12
13        with self.writer.as_default():
14            tf.summary.image("inputs", self.sample_images, step=epoch, max_outputs=4)
15            tf.summary.image("predictions", preds, step=epoch, max_outputs=4)
16
17# Example tensors
18sample_images = tf.random.uniform((4, 64, 64, 3), dtype=tf.float32)
19writer = tf.summary.create_file_writer("logs/image_demo/images")

Attach both callbacks to model training:

1# model.fit(
2#     train_ds,
3#     validation_data=val_ds,
4#     epochs=10,
5#     callbacks=[tb_callback, ImageSummaryCallback(sample_images, writer)],
6# )

Then run TensorBoard and open the Images tab.

Logging Grids and Preprocessing Views

For debugging preprocessing, log intermediate tensors such as normalized images, augmented samples, or segmentation masks. If masks are single channel, ensure shape still includes channel dimension.

1import tensorflow as tf
2
3def log_augmented_examples(writer, images, step):
4    augmented = tf.image.random_flip_left_right(images)
5    augmented = tf.image.random_brightness(augmented, max_delta=0.1)
6    augmented = tf.clip_by_value(augmented, 0.0, 1.0)
7
8    with writer.as_default():
9        tf.summary.image("augmented", augmented, step=step, max_outputs=4)

This makes it easy to spot pipeline issues such as wrong channel order, over aggressive augmentation, or unexpected normalization ranges.

Practical Workflow Tips

Use a fixed sample batch for epoch to epoch comparison, plus occasional random batches for coverage. Keep logged image counts modest to avoid huge event files. A few representative samples are usually enough.

Use separate tags for inputs, labels, and predictions so visual comparisons are clear. For segmentation or detection tasks, consider overlay rendering in preprocessing code before logging.

If training runs on remote machines, store logs in stable paths and sync them for local TensorBoard inspection.

For long experiments, split logs by run id and hyperparameter tag so image streams stay comparable. A clear directory convention such as model name, date, and seed makes debugging much easier when several experiments run in parallel.

It is also useful to log occasional failure cases separately, for example images with highest reconstruction error or misclassified samples. Reviewing these targeted images can reveal dataset issues that scalar loss curves cannot show.

Common Pitfalls

A common pitfall is logging tensors with wrong shape, such as missing batch dimension. tf.summary.image expects four dimensions.

Another issue is logging values outside display range. Tensors with large raw values can appear black or washed out. Normalize or clip before summary write.

Developers also log too many images every step, which inflates event files and slows training. Log at epoch boundaries or every few hundred steps.

Finally, using random samples only can make trend comparison difficult. Include a fixed reference batch to track qualitative improvements consistently.

Summary

• Use a custom Keras callback to write image summaries during training.
• Log tensors in batch height width channel shape with valid value ranges.
• Separate tags for inputs and predictions improve debugging clarity.
• Keep image logging frequency moderate to control log size.
• Combine fixed and random samples for both trend and coverage checks.