A modern GPU can train most vision and tabular models faster than a naive input pipeline can feed it. When a training job reports low GPU utilization, the model is rarely the problem — the pipeline is. This is the tf.data checklist we apply, in order, because the early items are cheap and frequently sufficient.
First, measure
Before touching the pipeline, confirm it is actually the bottleneck. The TensorBoard Profiler's input-pipeline analyzer states directly what fraction of each step is spent waiting on input. A cruder but useful probe: benchmark the dataset alone by iterating it without the model. If the pipeline alone cannot hit your target steps/second, no amount of model-side tuning matters.
import time
def benchmark(ds, steps=200):
it = iter(ds)
next(it) # warm-up
t0 = time.perf_counter()
for _ in range(steps):
next(it)
dt = time.perf_counter() - t0
print(f"{steps / dt:.1f} batches/sec")
1. prefetch at the end
The single highest-value line in tf.data:
ds = ds.prefetch(tf.data.AUTOTUNE)
Prefetching overlaps producer and consumer, so the pipeline builds the next batch while the accelerator works on the current one. It belongs at the very end of the chain, after batching. If a pipeline lacks it, add it before investigating anything else.
2. Parallelize map
Every nontrivial map should declare parallelism:
ds = ds.map(decode_and_augment, num_parallel_calls=tf.data.AUTOTUNE)
The default is sequential. AUTOTUNE lets the runtime allocate threads dynamically, which beats hand-tuned constants in almost every case we've measured. If op-level determinism doesn't matter for your job, deterministic=False on the map buys additional throughput by letting results arrive out of order.
3. cache what is expensive and stable
If the decoded dataset fits in RAM, ds.cache() after decoding but before augmentation makes every epoch after the first read from memory instead of redoing I/O and decode work. Two rules keep it safe:
- Cache before random augmentation, never after — caching after augmentation freezes one random draw and silently destroys the augmentation.
- For datasets larger than RAM,
cache(filename)spills to local disk, which still beats re-reading and re-decoding many small originals.
4. Fix the small-files problem at the source
Reading tens of thousands of small image files individually is the classic pipeline killer, especially from network or cloud storage where per-file latency dominates. The fix is sequential reads over a few large containers: pack the dataset into TFRecord shards (or another large-container format) sized in the hundreds of megabytes, then read them with parallel interleave:
ds = tf.data.Dataset.list_files("data/train-*.tfrecord")
ds = ds.interleave(
tf.data.TFRecordDataset,
cycle_length=tf.data.AUTOTUNE,
num_parallel_calls=tf.data.AUTOTUNE,
)
This is a preprocessing-job change, not a pipeline flag, which is why teams defer it — and why it is usually the biggest single win on image workloads.
5. Vectorize the map when you can
Per-element Python-level work is expensive. When a transformation is expressed in vectorized ops, batching before mapping amortizes op dispatch across the whole batch:
ds = ds.batch(batch_size).map(augment_batch, num_parallel_calls=tf.data.AUTOTUNE)
Related: avoid tf.py_function in hot paths. It punches out to the Python interpreter, serializing on the GIL and defeating graph-level parallelism. Most uses can be rewritten with native TF ops; the stubborn remainder (bespoke scientific decoders, for instance) should run in the offline preprocessing job instead.
6. Get shuffle right, then stop worrying
shuffle(buffer_size) trades memory for randomization quality. A buffer far smaller than the dataset yields locally-correlated batches, which can quietly hurt convergence — a modeling bug caused by a pipeline setting. Practical guidance: shuffle filenames fully (that's cheap), give the record-level shuffle as large a buffer as memory allows, and place shuffle before batching. If you sharded thoughtfully at creation time, a moderate record buffer on top of full filename shuffling is fine.
7. Order the pipeline deliberately
A shape that serves most training jobs:
list_files → shuffle(files) → interleave(read) → map(decode)
→ cache → shuffle(records) → map(augment) → batch → prefetch
Expensive-and-deterministic work sits before cache; random work sits after it; prefetch is last. Deviations should be conscious decisions, not accidents of code history.
What we see in practice
On review engagements, the most common findings are, in order: no prefetch, sequential map, thousands of small files read directly from object storage, and augmentation accidentally placed before cache. The first two are one-line fixes. The third is an afternoon of preprocessing work that routinely doubles throughput on image workloads. None of them require touching the model — which is exactly why they're worth checking before any conversation about architecture or hardware.