Module 5 — tf.data: reading, transforming and prefetching

An idle accelerator costs exactly as much as a busy one. Yet in a poorly fed training run it spends most of its time waiting for data. This module addresses that bottleneck, which is both the most common and the most profitable to fix.

The problem: accelerator starvation

Without a pipeline, everything is strictly sequential. The processor reads a batch, transforms it, hands it over; meanwhile the accelerator waits. The accelerator computes; meanwhile the processor waits. The two never work at the same time, and total time is the sum of both.

tf.data builds a pipeline where reading and computing overlap. Total time becomes the maximum of the two instead of their sum. On a run where data preparation takes 40 % of the time, the gain is immediate and substantial.

Composing a pipeline

import tensorflow as tf

AUTO = tf.data.AUTOTUNE

dataset = (
    tf.data.Dataset.from_tensor_slices((paths, labels))
    .shuffle(10_000)
    .map(load_and_resize, num_parallel_calls=AUTO)
    .batch(32)
    .prefetch(AUTO)
)

model.fit(dataset, epochs=20)

Each link has a precise role:

Operation	Role	Note
from_tensor_slices	splits into elements	keep paths, not images
shuffle(n)	shuffles within a buffer of n	too small shuffles nothing
map(f)	applies a transformation	parallelise with num_parallel_calls
batch(k)	groups into batches of k	after shuffling, never before
prefetch	prepares the next batch	always last

AUTOTUNE lets TensorFlow choose the parallelism level and buffer depth by measuring at runtime. It is almost always better than a hand-picked value, and it adapts to the machine.

Operation order changes the outcome

This is not a matter of style: several permutations produce wrong or slow results.

shuffle before batch. Shuffling after grouping only reorders the batches, whose contents stay identical from one epoch to the next. The model always sees the same neighbourhoods, and the regularising effect of shuffling disappears.

batch before map when the transformation vectorises. Applying a normalisation element by element costs one function call per example. On a batch, it is a single call on a tensor. For purely arithmetic operations, swapping the order is noticeably faster. Image decoding, on the other hand, must stay per element, since each file is distinct.

prefetch last, always. Placed in the middle, it only overlaps the tail of the pipeline and leaves the accelerator waiting on the upstream stages.

The shuffle buffer is not the dataset

shuffle(1000) fills a thousand-element buffer and draws from it. On a hundred-thousand-example dataset sorted by class, that buffer holds a single class at a time: the shuffling is an illusion and every batch is single-class. The model diverges for no visible reason. Use a buffer on the order of the dataset size, or shuffle the file paths upstream with shuffle on a Python list, which costs no memory at all.

Where to place the cache

cache memorises the result of everything before it, so its position determines what gets reused.

dataset = (
    tf.data.Dataset.from_tensor_slices((paths, labels))
    .map(decode_image, num_parallel_calls=AUTO)
    .cache()                       # after decoding, before augmentation
    .shuffle(10_000)
    .map(augment, num_parallel_calls=AUTO)
    .batch(32)
    .prefetch(AUTO)
)

Decoding always yields the same result: caching it saves identical work on every epoch. Augmentation, by contrast, must produce fresh variation on each pass; placing it before the cache would freeze a single augmented version per image and defeat its entire purpose.

An in-memory cache is only viable if the decoded dataset fits. Otherwise cache("/path/to/file") writes to disk, which is still far faster than repeated decoding.

Reading from files

For images on disk, image_dataset_from_directory covers the common case in one line, inferring classes from folder names:

dataset = keras.utils.image_dataset_from_directory(
    "data/train",
    image_size=(224, 224),
    batch_size=32,
    label_mode="int",
)

At larger volumes the TFRecord format becomes preferable. It groups examples into a few large sequential files, which removes the cost of opening millions of small files — often the dominant factor on network storage.

dataset = (
    tf.data.TFRecordDataset(files, num_parallel_reads=AUTO)
    .map(parse_example, num_parallel_calls=AUTO)
    .batch(64)
    .prefetch(AUTO)
)

The Python function trap

A function passed to map is traced, as in module 1. It must therefore be expressed in TensorFlow operations. An arbitrary Python library has no place there.

# does not work: PIL is not a TensorFlow operation
def load(path):
    return numpy.array(PIL.Image.open(path.numpy()))

Two ways out. The good one uses native operations, tf.io.read_file then tf.image.decode_jpeg. The other wraps the Python code in tf.py_function, which works but serialises execution on the global interpreter lock: parallelism disappears, and with it most of the benefit.

Measure before optimising

Time one epoch with your full pipeline, then with dataset.take(1).repeat(), which replays a batch already in memory. If the second is markedly faster, data is the bottleneck and this module applies. If both times are close, computation dominates and you should look elsewhere — model size, mixed precision, or distribution in module 9.

Key takeaways

• Without a pipeline, total time is the sum of preparation and computation; with prefetch it becomes their maximum.
• Order is binding: shuffle before batch, batch before map for vectorisable transformations, and prefetch always last.
• The shuffle buffer must be comparable to the dataset size, otherwise a sorted dataset yields single-class batches and training diverges for no visible reason.
• cache belongs after deterministic transformations and before augmentation, or augmentation would freeze on a single version per example.

Next module: callbacks, which let you act during training without rewriting it.