DeepAdvisors | tf.data Pipelines That Keep Your GPU Busy

The most common performance problem we find on client training jobs is not the model. It is a GPU sitting at 30% utilization waiting for the next batch. Input pipelines are where training time goes to die, and tf.data gives you everything you need to fix it - if you use it deliberately. This is the checklist we run on every training job we touch.

First, prove it is the input pipeline

Do not optimize on a hunch. The TensorFlow Profiler will tell you whether a step is input-bound. Capture a handful of steps with the TensorBoard callback:

import keras

tb = keras.callbacks.TensorBoard(log_dir="logs", profile_batch=(50, 60))
model.fit(train_ds, epochs=1, callbacks=[tb])

Open TensorBoard, go to the Profile tab, and read the Overview Page. It reports the percentage of step time spent waiting on input and says, in words, whether your program is input-bound. The Input Pipeline Analyzer then breaks down which tf.data op is the bottleneck. The Profiler guide covers the rest of the tool.

If the overview says you are not input-bound, stop here and go look at the model. If it says you are, continue.

The baseline pipeline

Here is the shape of a pipeline we see constantly - correct, and slow:

import tensorflow as tf

def parse(example):
    features = tf.io.parse_single_example(example, {
        "image": tf.io.FixedLenFeature([], tf.string),
        "label": tf.io.FixedLenFeature([], tf.int64),
    })
    image = tf.io.decode_jpeg(features["image"], channels=3)
    image = tf.image.resize(image, [224, 224]) / 255.0
    return image, features["label"]

files = tf.io.gfile.glob("data/train-*.tfrecord")
ds = (
    tf.data.TFRecordDataset(files)
    .map(parse)
    .shuffle(1000)
    .batch(64)
)

Everything runs sequentially on one thread, and the model waits for each batch to be built. Let us fix it one transformation at a time.

1. `prefetch`: overlap input with training

The single highest-value change. prefetch lets the pipeline prepare the next batch while the current one trains:

ds = ds.prefetch(tf.data.AUTOTUNE)

Always the last call in the chain. AUTOTUNE lets the runtime pick the buffer size; you rarely need to override it.

2. Parallelize `map`

Decoding and resizing images is CPU work that parallelizes well:

ds = ds.map(parse, num_parallel_calls=tf.data.AUTOTUNE)

If your map function is cheap (a few tensor ops on small inputs), the per-element scheduling overhead can outweigh the gain; in that case batch first and write a vectorized map that operates on the whole batch.

3. `interleave`: read files in parallel

TFRecordDataset(files) reads the shards one after another. interleave reads from several at once, which matters enormously on network filesystems and object storage:

ds = tf.data.Dataset.from_tensor_slices(files).shuffle(len(files))
ds = ds.interleave(
    tf.data.TFRecordDataset,
    cycle_length=8,
    num_parallel_calls=tf.data.AUTOTUNE,
    deterministic=False,
)

deterministic=False allows elements to arrive out of order for throughput; leave it at the default True only if you need reproducible ordering.

4. `cache`: do expensive work once

If the dataset fits in memory after preprocessing, cache it after the expensive transformations and before the per-epoch ones (shuffle, augmentation):

ds = ds.map(parse, num_parallel_calls=tf.data.AUTOTUNE).cache()
ds = ds.shuffle(10_000).map(augment, num_parallel_calls=tf.data.AUTOTUNE)
ds = ds.batch(64).prefetch(tf.data.AUTOTUNE)

cache("path/on/disk") spills to local disk for datasets that do not fit in RAM, which is still usually faster than re-reading from remote storage. Two warnings: never cache after random augmentation (every epoch will see identical samples), and the first epoch pays the full cost, so do not benchmark on it.

5. Order matters

A good ordering for image training:

interleave file reads
map(parse) in parallel
cache (if it fits)
shuffle with a buffer large enough to actually mix the data
map(augment) in parallel
batch
prefetch

Put batch before map when the map function is vectorizable; put shuffle before batch always (shuffling batches is not shuffling data).

The optimized pipeline

AUTOTUNE = tf.data.AUTOTUNE

files = tf.io.gfile.glob("data/train-*.tfrecord")
ds = (
    tf.data.Dataset.from_tensor_slices(files)
    .shuffle(len(files))
    .interleave(
        tf.data.TFRecordDataset,
        cycle_length=8,
        num_parallel_calls=AUTOTUNE,
        deterministic=False,
    )
    .map(parse, num_parallel_calls=AUTOTUNE)
    .cache()
    .shuffle(10_000)
    .map(augment, num_parallel_calls=AUTOTUNE)
    .batch(64, drop_remainder=True)
    .prefetch(AUTOTUNE)
)

Re-profile after the change. On a typical image-classification job we see input-bound step time drop from the majority of the step to a few percent, and GPU utilization climb accordingly. Your numbers will differ; the profiler is the source of truth, not this post.

TFRecord vs Parquet in 2026

Clients ask this constantly, because their data engineering teams live in Parquet.

TFRecord is still the fastest path into tf.data: it is a simple sequential format, TFRecordDataset is native, and interleave over shards scales cleanly. The cost is an export step and a format nobody outside ML uses.

Parquet is the lingua franca of the data platform: columnar, compressed, queryable, and already where your data lives. Reading it into tf.data means going through Arrow (for example pyarrow in a from_generator pipeline or a Python-level reader), which adds a Python hop and is harder to parallelize with AUTOTUNE.

Our rule of thumb:

Small or medium tabular datasets: read Parquet with Arrow into memory, build the dataset with from_tensor_slices, cache. No TFRecord needed.
Large image, audio, or sequence datasets trained repeatedly: export to sharded TFRecords once (a few hundred MB per shard), then the pipeline above. The export pays for itself on the second training run.
Data that changes daily: make the TFRecord export part of the data platform's job, not a notebook someone runs.

Other things worth knowing

tf.data.Dataset.snapshot materializes a preprocessed pipeline to disk for reuse across runs - a heavier alternative to cache when preprocessing is very expensive.
options.experimental_optimization lets you toggle autotuning and fusion behaviours; the defaults are good and you should change them only with a profile in hand.
tf.data works on every Keras 3 backend. A tf.data pipeline feeds model.fit whether the model runs on TensorFlow, JAX, or PyTorch, so this work is not wasted if you change backends.
Distributed training multiplies the problem: each replica needs its own fast input stream. Shard at the file level (ds.shard or interleave over a per-worker file subset) and prefetch per replica.

The full reference is the tf.data performance guide; this post is the subset that fixes most real jobs.

If your training runs are slower than they should be and you are not sure why, a profiling session is usually a one-day engagement. Contact us and we will start with the profiler, not a rewrite.

tf.data Pipelines That Keep Your GPU Busy