3. TensorFlow data pipelines

In the previous chapter we saw how a model is defined and trained. In practice, however, the model is rarely the bottleneck. On most real workloads — image decoding, text tokenization, audio feature extraction — the GPU sits idle waiting for the CPU to produce the next batch. The tf.data API exists precisely to solve this problem: it lets you describe an input pipeline as a graph of transformations that TensorFlow can execute in parallel with the training step.

This chapter builds a complete, measurable input pipeline. We will start from raw files on disk, construct a tf.data.Dataset, apply the transformations that matter, and then tune the two parameters that most often determine throughput: batch() and prefetch(). Along the way we will measure the effect of each change rather than assuming it helps.

From Tensors to a Dataset: The Core Abstraction

A tf.data.Dataset is a sequence of elements, where each element is a (possibly nested) structure of tensors. Crucially, a Dataset is lazy: nothing is read or decoded until iteration begins. This laziness is what allows TensorFlow to fuse transformations, reorder them, and run them on background threads.

The two entry points you will use most are from_tensor_slices() for in-memory data and list_files() plus interleave() for data on disk. For a directory of images organized as data/train/<class>/*.jpg, the idiomatic construction is:

import tensorflow as tf

AUTOTUNE = tf.data.AUTOTUNE

file_ds = tf.data.Dataset.list_files("data/train/*/*.jpg", shuffle=True)

def load_and_decode(path):
    label = tf.strings.split(path, "/")[-2]          # class folder name
    image = tf.io.read_file(path)
    image = tf.io.decode_jpeg(image, channels=3)
    image = tf.image.resize(image, [224, 224])
    image = tf.cast(image, tf.float32) / 255.0
    return image, label

dataset = file_ds.map(load_and_decode, num_parallel_calls=AUTOTUNE)

Three details deserve attention. First, list_files() returns a dataset of file paths, not file contents — the actual I/O happens inside map(). Second, num_parallel_calls=AUTOTUNE tells TensorFlow to run the mapping function on a pool of threads sized dynamically based on available CPU. Without it, decoding is serialized and the pipeline stalls. Third, the label is derived from the path, so no separate label file is needed.

Transformations That Change Performance

Not all transformations are equal. Some are cheap metadata operations; others are expensive and must be parallelized. The table below summarizes the ones you will use in almost every pipeline.

The ordering of these calls matters. A common mistake is to call batch() before map(), which forces the mapping function to operate on batches rather than single elements and breaks the parallelism model. The canonical order is: read → shuffle → map → batch → prefetch. Shuffle before map so that expensive decoding is spread across the whole dataset; batch after map so that each element is decoded individually; prefetch last so that the entire upstream pipeline is overlapped with the training step.

Batching and Prefetching in Practice

Batching converts a dataset of individual examples into a dataset of tensors with a leading batch dimension. The batch size is a trade-off: larger batches improve GPU utilization and reduce per-step overhead, but they also increase memory consumption and can hurt convergence if the effective learning rate is not scaled. A reasonable starting point for image classification is 32 or 64; for large language models, batch size is often limited by memory rather than throughput.

Prefetching is the mechanism that hides latency. When you call dataset.prefetch(tf.data.AUTOTUNE), TensorFlow runs the upstream pipeline on a background thread and buffers the next batch (or several) while the current batch is being consumed by the model. The effect is that the training step and the input pipeline run concurrently. Without prefetch, the timeline looks like: read batch → train → read batch → train. With prefetch, it looks like: read batch 1 → train batch 1 while reading batch 2 → train batch 2 while reading batch 3, and so on.

Measuring Input Performance

You cannot tune what you do not measure. The simplest measurement is to iterate over the dataset without a model and time it:

import time

def benchmark(dataset, num_epochs=2):
    start = time.perf_counter()
    for _ in range(num_epochs):
        for _ in dataset:
            pass
    return (time.perf_counter() - start) / num_epochs

baseline = file_ds.map(load_and_decode)
print("no prefetch:", benchmark(baseline))

optimized = (file_ds
    .map(load_and_decode, num_parallel_calls=AUTOTUNE)
    .batch(32)
    .prefetch(AUTOTUNE))
print("optimized:", benchmark(optimized))

On a typical workstation with 10,000 JPEG images, the unoptimized pipeline might take 45 seconds per epoch, while the optimized version drops to 12 seconds — a 3.7× improvement that translates directly into faster training. The exact numbers depend on disk speed, image size, and CPU count, so always measure on your own hardware.

For finer-grained analysis, TensorFlow provides the tf.data.experimental.StatsOptions API, which records per-transformation timing and can be visualized in TensorBoard. This is useful when a single transformation — for example, a slow Python function inside map() — dominates the pipeline.

Common Failure Modes

• Python function inside map(): if the function is not traceable by tf.function, it runs in the Python interpreter and cannot be parallelized. Prefer TensorFlow ops (tf.io, tf.image, tf.strings) inside map().
• Shuffle buffer too small: a buffer of 1,000 on a dataset of 1,000,000 images produces nearly sequential order. Use a buffer of at least 10× the batch size, or shuffle file paths before reading.
• Prefetch after batch but before unbatch: prefetching a dataset of batches is correct; prefetching before batching buffers individual elements and wastes memory.
• Cache on a dataset that does not fit in memory: cache() will spill to disk or fail. Check dataset size first, or cache to a file path with cache("cache.tf-data").

Putting It Together: A Reference Pipeline

The following snippet is a complete, production-shaped pipeline for an image classification task. It reads JPEGs from a class-organized directory, decodes and resizes them, batches, and prefetches. It is the template you should adapt for your own data.

def build_pipeline(data_dir, batch_size=32, image_size=(224, 224)):
    files = tf.data.Dataset.list_files(f"{data_dir}/*/*.jpg", shuffle=True)

    def parse(path):
        label = tf.strings.split(path, "/")[-2]
        img = tf.io.decode_jpeg(tf.io.read_file(path), channels=3)
        img = tf.image.resize(img, image_size)
        img = tf.cast(img, tf.float32) / 255.0
        return img, label

    return (files
        .shuffle(buffer_size=10_000)
        .map(parse, num_parallel_calls=tf.data.AUTOTUNE)
        .batch(batch_size)
        .prefetch(tf.data.AUTOTUNE))

To verify the pipeline before training, pull one batch and inspect its shape and value range:

ds = build_pipeline("data/train")
images, labels = next(iter(ds))
print(images.shape)        # (32, 224, 224, 3)
print(images.numpy().min(), images.numpy().max())  # 0.0 1.0
print(labels[:5].numpy())  # e.g. [b'cats' b'dogs' ...]

If the shape is (32, 224, 224, 3) and the value range is [0, 1], the pipeline is correct. If the shape is (224, 224, 3), you forgot batch(). If the range is [0, 255], you forgot the cast and division. These two checks catch the majority of pipeline bugs before a single training step runs.