Your GPU Is Bored: Fixing Input Pipelines with tf.data

The expensive hardware that isn't doing anything

You bought the GPU. You wrote the model. Training starts, and your steps-per-second number is a fraction of what the hardware should deliver. Before you blame the model architecture, check the boring part: the input pipeline. In a surprising number of TensorFlow training jobs, the GPU spends a large share of each step idle, waiting for the CPU to finish decoding images, tokenizing text, or parsing records. This is called GPU starvation, and it's usually fixable in a dozen lines of code.

The thesis of this post: most input-pipeline bottlenecks come down to ordering and parallelism in tf.data, and the fix is a small set of transformations applied in the right sequence — with prefetch last.

How tf.data thinks about work

The tf.data API builds pipelines as a functional chain: you start with a source (files, tensors, a generator) and attach transformations that each produce a new dataset. Two kinds of work happen here. The producer side runs on the CPU — reading files, decoding, augmenting. The consumer is your training step on the GPU. If the producer can't keep up, the consumer sits idle between batches.

Four transformations do most of the heavy lifting:

• interleave() reads from multiple files in parallel, which maximizes disk throughput when your data is sharded across many TFRecord files.
• map() applies your preprocessing function, and with num_parallel_calls it fans that work out across CPU cores.
• cache() stores the dataset (in memory, or in a file if you pass a filename) after the first epoch, so epochs two onward skip upstream work entirely.
• prefetch() overlaps producer and consumer: while the GPU trains on batch n, the CPU prepares batch n+1.

A worked example: the order matters

Here's a typical image-training pipeline, written for TensorFlow 2.x (behavior described here is stable across recent 2.x releases, but confirm against your installed version):

import tensorflow as tf

files = tf.data.Dataset.list_files("/data/train/*.tfrecord", shuffle=True)

ds = (
    files
    .interleave(tf.data.TFRecordDataset,
                cycle_length=tf.data.AUTOTUNE,
                num_parallel_calls=tf.data.AUTOTUNE)
    .shuffle(buffer_size=10_000)
    .map(parse_and_augment, num_parallel_calls=tf.data.AUTOTUNE)
    .batch(128)
    .prefetch(tf.data.AUTOTUNE)
)

Run this in your training script wherever you currently construct the dataset; it needs no special permissions beyond read access to the data. A few deliberate choices:

• interleave comes first so file reading is parallel from the start.
• map gets num_parallel_calls=tf.data.AUTOTUNE, letting the runtime pick a parallelism level based on available cores instead of you guessing a number.
• prefetch is last. This is the part people get wrong. Prefetch overlaps everything upstream of it with the training step. If you prefetch early — say, before a heavy map — you only overlap cheap file reads, and the expensive augmentation still runs serially between steps.

Where does cache() fit? Place it after expensive, deterministic work you don't want to repeat — for example, after decoding but before random augmentation, so you re-decode zero times but still get fresh augmentations each epoch:

ds = ds.map(decode_only).cache().map(random_augment, num_parallel_calls=tf.data.AUTOTUNE)

Prove it worked — and know the trade-offs

Don't take the reordering on faith. Two concrete checks:

1. Steps per second. Log step time before and after the change. This is the cheapest signal and often the only one you need.
2. TensorFlow Profiler. Capture a profile (via the TensorBoard callback with profile_batch) and look at the step-time breakdown. If a large slice is labeled as input/wait time, the pipeline is still the bottleneck; if the GPU kernel time dominates, you've fixed it and further pipeline tuning is wasted effort.

Now the honest limitations. cache() holds the entire cached dataset in host RAM if you don't give it a file path — for a large image dataset that's a straightforward route to an out-of-memory kill on the machine, not the GPU. Watch system RAM during the first epoch; if it climbs toward the limit, cache to a file on fast local disk instead, or drop caching. And more parallelism isn't free: on a shared or small machine, aggressive num_parallel_calls can starve the very Python process driving training. AUTOTUNE helps, but it optimizes for throughput, not for being a good neighbor.

What to do on Monday

Take your slowest training job, move prefetch(tf.data.AUTOTUNE) to the end of the pipeline, add num_parallel_calls=tf.data.AUTOTUNE to your map, and measure steps per second again. If the GPU is still idle, profile before adding more machinery — the bottleneck may have moved to the model itself, which is exactly where you want it.