TFRecord
TensorFlow
image processing
data pipeline
machine learning

How to read images with different size in a TFRecord file

Master System Design with Codemia

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

Introduction

TFRecord can store images of different sizes without any problem, but your input pipeline has to decode them correctly and then decide how to batch them. The main trap is assuming that all images have one fixed shape at parse time. In practice, you usually store encoded image bytes, decode each example individually, and only resize or pad after decoding.

Store Encoded Bytes, Not Raw Fixed-Shape Tensors

The usual pattern is to serialize each image as encoded bytes, plus optional metadata such as label, height, and width.

Example writer:

python
1import tensorflow as tf
2
3
4def _bytes_feature(value):
5    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))
6
7
8def _int64_feature(value):
9    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))
10
11
12with tf.io.TFRecordWriter("images.tfrecord") as writer:
13    for image_path, label in [("a.jpg", 0), ("b.jpg", 1)]:
14        image_bytes = tf.io.read_file(image_path).numpy()
15        image = tf.io.decode_image(image_bytes, channels=3)
16
17        example = tf.train.Example(features=tf.train.Features(feature={
18            "image_bytes": _bytes_feature(image_bytes),
19            "height": _int64_feature(int(image.shape[0])),
20            "width": _int64_feature(int(image.shape[1])),
21            "label": _int64_feature(label),
22        }))
23
24        writer.write(example.SerializeToString())

The variable-sized part is stored safely inside image_bytes. That is what makes heterogeneous image shapes practical in one TFRecord file.

Parse and Decode Each Example First

When reading, parse the serialized example and decode the image bytes:

python
1import tensorflow as tf
2
3feature_spec = {
4    "image_bytes": tf.io.FixedLenFeature([], tf.string),
5    "height": tf.io.FixedLenFeature([], tf.int64),
6    "width": tf.io.FixedLenFeature([], tf.int64),
7    "label": tf.io.FixedLenFeature([], tf.int64),
8}
9
10
11def parse_example(serialized):
12    parsed = tf.io.parse_single_example(serialized, feature_spec)
13    image = tf.io.decode_jpeg(parsed["image_bytes"], channels=3)
14    label = parsed["label"]
15    return image, label
16
17
18dataset = tf.data.TFRecordDataset(["images.tfrecord"]).map(parse_example)
19
20for image, label in dataset.take(2):
21    print(image.shape, label.numpy())

At this stage, TensorFlow knows the tensor rank but not a fixed height and width shared by the whole dataset. That is expected.

Resize or Pad Before Regular Batching

Regular batch() requires compatible tensor shapes. If your images are different sizes, you usually do one of these:

  • resize all images to a common target
  • pad them to the largest shape in a batch
  • keep them unbatched for a custom consumer

The simplest path is resize:

python
1def parse_and_resize(serialized):
2    parsed = tf.io.parse_single_example(serialized, feature_spec)
3    image = tf.io.decode_jpeg(parsed["image_bytes"], channels=3)
4    image = tf.image.resize(image, [224, 224])
5    label = parsed["label"]
6    return image, label
7
8
9dataset = (
10    tf.data.TFRecordDataset(["images.tfrecord"])
11    .map(parse_and_resize)
12    .batch(32)
13    .prefetch(tf.data.AUTOTUNE)
14)

That is the standard training pipeline approach because most models require a fixed input size anyway.

Use padded_batch When You Need Variable Original Shapes

If you want to preserve more of the original size information, padded_batch is a better fit:

python
1dataset = tf.data.TFRecordDataset(["images.tfrecord"]).map(parse_example)
2
3dataset = dataset.padded_batch(
4    2,
5    padded_shapes=([None, None, 3], []),
6    padding_values=(tf.constant(0, dtype=tf.uint8), tf.constant(0, dtype=tf.int64))
7)
8
9for images, labels in dataset.take(1):
10    print(images.shape)
11    print(labels.numpy())

This pads each image in the batch to a common height and width for that batch only. It is useful when resizing would damage the task or when the downstream model supports masking or variable spatial dimensions.

Keep Shape Decisions Close to the Model

The right strategy depends on the model and task:

  • classification usually resizes to a fixed input shape
  • detection and segmentation may prefer padding or smarter aspect-ratio handling
  • archival or inspection pipelines may keep original shapes as long as possible

Do not force one shape policy into the TFRecord design itself. The TFRecord should store decodable examples cleanly; the input pipeline should decide how to normalize them for the model.

Common Pitfalls

  • Storing raw pixel tensors and assuming all examples can share one fixed shape forever.
  • Trying to use ordinary batch() before resizing or padding variable-sized images.
  • Forgetting that decode_jpeg returns dynamic spatial dimensions until later pipeline steps constrain them.
  • Resizing blindly without checking whether the task actually tolerates geometric distortion.
  • Treating TFRecord as the place to solve batching strategy instead of the dataset pipeline.

Summary

  • TFRecord can store images of different sizes by keeping them as encoded bytes.
  • Parse and decode each example first, then decide how to normalize shapes.
  • Use resizing for fixed-shape model input pipelines.
  • Use padded_batch when preserving variable dimensions is more important.
  • Keep storage format and batching strategy as separate decisions in the data pipeline.

Course illustration
Course illustration

All Rights Reserved.