TensorBoard
image visualization
deep learning
machine learning
Python

How to plot grid of images in tensorboard?

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

TensorBoard can display images directly with tf.summary.image, but if you want a true grid, you need to arrange the batch into one tiled image first. That usually means reshaping and concatenating the image tensor so several samples appear in a single canvas. Once you have that tiled tensor, you log it as one image summary.

What TensorBoard Expects

tf.summary.image expects a tensor shaped like:

text
[N, H, W, C]

where:

  • 'N is the number of images'
  • 'H is height'
  • 'W is width'
  • 'C is channels, typically 1, 3, or 4'

If you pass a batch directly, TensorBoard will show multiple separate images up to max_outputs. That is useful, but it is not the same as a manually arranged grid.

Build a Grid Tensor First

A common approach is to tile rows * cols images into one larger image.

python
1import tensorflow as tf
2
3def image_grid(images, rows, cols):
4    images = tf.convert_to_tensor(images)
5    images = images[:rows * cols]
6
7    batch, height, width, channels = images.shape
8    if batch < rows * cols:
9        raise ValueError("Not enough images to fill the grid")
10
11    images = tf.reshape(images, [rows, cols, height, width, channels])
12    images = tf.transpose(images, [0, 2, 1, 3, 4])
13    images = tf.reshape(images, [1, rows * height, cols * width, channels])
14    return images
15
16sample_images = tf.random.uniform([16, 28, 28, 1])
17grid = image_grid(sample_images, rows=4, cols=4)
18print(grid.shape)

The returned tensor has shape [1, big_height, big_width, channels], which is exactly what tf.summary.image needs for logging one composite image.

Log the Grid to TensorBoard

python
1import tensorflow as tf
2
3log_dir = "logs/images"
4writer = tf.summary.create_file_writer(log_dir)
5
6sample_images = tf.random.uniform([16, 28, 28, 1])
7grid = image_grid(sample_images, rows=4, cols=4)
8
9with writer.as_default():
10    tf.summary.image("sample_grid", grid, step=0)
11    writer.flush()

Then run:

bash
tensorboard --logdir logs/images

Open TensorBoard and look under the Images tab.

Normalization and Data Range

TensorBoard expects image values in a reasonable display range. If your images are floating-point tensors, a [0, 1] range is usually the safest choice.

For example:

python
images = tf.cast(images, tf.float32)
images = tf.clip_by_value(images, 0.0, 1.0)

If your data comes in [-1, 1], rescale it first:

python
images = (images + 1.0) / 2.0
images = tf.clip_by_value(images, 0.0, 1.0)

Otherwise the grid may appear washed out, too dark, or inconsistent.

Use It During Training

A common pattern is to log image grids periodically inside a training loop so you can inspect:

  • input batches
  • model reconstructions
  • generated images
  • segmentation predictions

Example inside an epoch loop:

python
1for step in range(3):
2    images = tf.random.uniform([16, 28, 28, 1])
3    grid = image_grid(images, 4, 4)
4    with writer.as_default():
5        tf.summary.image("training_samples", grid, step=step)

This gives you a timeline of visual changes across training steps.

A Note on Alternatives

If you already use PyTorch or Matplotlib-style tooling, you may know helper utilities that create grids automatically. In TensorFlow-only workflows, the explicit tensor-reshaping approach is often simplest because it avoids extra dependencies and keeps the data inside the TensorFlow pipeline.

That also makes it easier to log model outputs without converting back and forth between formats.

Common Pitfalls

The most common mistake is assuming max_outputs=16 automatically creates a 4 by 4 grid. It does not. It logs up to 16 separate images.

Another mistake is passing tensors with the wrong shape, such as [H, W, C] instead of [N, H, W, C]. tf.summary.image expects a batch dimension.

Developers also often forget to normalize image values, which makes the logged result look wrong even when the tensor layout is correct.

Finally, if the grid appears scrambled, check the reshape and transpose order. A wrong transpose usually causes row and column mixing.

Summary

  • 'tf.summary.image logs batches of images, but a real grid requires tiling first.'
  • Build the grid by reshaping, transposing, and flattening the image batch into one larger image.
  • Log the tiled tensor as a single image summary.
  • Normalize image values so TensorBoard displays them correctly.
  • Use the same pattern to visualize inputs, reconstructions, or generated samples during training.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.