The tf.data pipeline in depth

Build input pipelines that keep the GPU fed: map, batch, shuffle, cache, prefetch, and how to prove the bottleneck is the input rather than the model.

Building a pipeline

import tensorflow as tf

AUTOTUNE = tf.data.AUTOTUNE

def build_dataset(files, batch_size=64, shuffle=True):
    ds = tf.data.Dataset.from_tensor_slices(files)
    if shuffle:
        ds = ds.shuffle(buffer_size=len(files), seed=0, reshuffle_each_iteration=True)

    ds = ds.map(load_and_decode, num_parallel_calls=AUTOTUNE)
    ds = ds.cache()                            # in RAM after the first epoch
    ds = ds.map(train_augment, num_parallel_calls=AUTOTUNE)
    ds = ds.batch(batch_size, drop_remainder=True)
    ds = ds.map(normalise, num_parallel_calls=AUTOTUNE)
    ds = ds.prefetch(AUTOTUNE)                 # overlaps step n+1 with training step n
    return ds

train_ds = build_dataset(train_files)
val_ds = build_dataset(val_files, shuffle=False).take(20)

model.fit(train_ds, epochs=10, validation_data=val_ds)

cache() before augmentation and batch(): the raw decoded examples are reused each epoch, so the expensive decode happens once.
Cache after decoding and before augmentation. Caching after augmentation defeats the point — you would freeze one augmented version for every epoch.
Only cache() if the dataset fits in memory or on disk. For a large dataset, cache to a file path instead of the default in-memory cache.
prefetch(AUTOTUNE) is the single highest-value line. It lets the CPU prepare batch n+1 while the GPU trains on batch n.
Set reshuffle_each_iteration=True so epoch two sees a different order; otherwise shuffle only shuffles once.

Generators, ragged data and lookup tables

# a generator-backed dataset for data that does not fit in memory
def gen():
    with open("corpus.txt", encoding="utf-8") as fh:
        for line in fh:
            text, label = line.rstrip().split("	")
            yield text, int(label)

ds = tf.data.Dataset.from_generator(
    gen,
    output_signature=(tf.TensorSpec(shape=(), dtype=tf.string),
                      tf.TensorSpec(shape=(), dtype=tf.int32)),
).batch(32).prefetch(AUTOTUNE)

# a vocabulary as a lookup table, for categorical ids
lookup = tf.keras.layers.StringLookup(vocabulary=vocab, num_oov_indices=1)
ids = lookup(tf.constant(["red", "green", "unknown-colour"]))
# [1, 2, 0] — index 0 is reserved for out-of-vocabulary

# build the vocabulary from the data itself
lookup = tf.keras.layers.StringLookup()
lookup.adapt(train_ds.map(lambda text, _: tf.strings.split(text)))

output_signature is mandatory and must match exactly what the generator yields; type or rank mismatches fail at graph build time with a message about the signature.
Use from_generator only when nothing else fits: it runs Python in the input path, cannot be parallelised as freely, and breaks graph serialisation.
adapt() on a preprocessing layer learns the vocabulary from data and produces a layer you can embed in the model, so the same transformation is applied at serving.
For ragged sequences, use tf.ragged or pad explicitly with padded_batch — a plain batch over variable-length tensors raises.

Finding the real bottleneck

import time

def time_pipeline(ds, n=20):
    it = iter(ds)
    next(it)                        # warm up: fill the prefetch buffer
    start = time.perf_counter()
    for _ in range(n):
        next(it)
    return (time.perf_counter() - start) / n

per_batch = time_pipeline(train_ds)
print(f"input pipeline: {per_batch * 1000:.1f} ms per batch")

# compare against pure model time for the same batch
batch = next(iter(train_ds))
start = time.perf_counter()
for _ in range(20):
    model(batch[0], training=True)
print(f"model step:   {(time.perf_counter() - start) / 20 * 1000:.1f} ms per batch")

Symptom	Likely cause	Fix
GPU utilisation under 60%	Input-bound pipeline	Add `prefetch`, `num_parallel_calls`
First epoch slow, later fast	Decoding without a cache	Put `cache()` after decode
Time per step grows each epoch	Growing in-memory cache	Cache to a file, or drop the cache
Step time far exceeds input time	Model is the bottleneck	Optimise the model, not the data
Memory blows up on shuffle	Buffer size equal to the dataset	Shuffle a window, not the whole set

Measure before optimising. If the model step time is larger than the pipeline time, no amount of pipeline tuning will help, and a profiler (tf.profiler, or model.fit with the TensorBoard callback) will show whether the time is in a specific kernel or in the input.

💡

Determinism and speed pull in opposite directions. num_parallel_calls with autotune makes example order depend on thread scheduling, so a reproducibility-sensitive run should fix the parallelism or seed the shuffle and accept the loss in throughput.

FAQ

How large should the shuffle buffer be?

At least a few thousand examples, and ideally the whole dataset if it fits. A buffer smaller than the batch size means no shuffling happens within a batch at all, which is worse than it sounds for ordered data.

Should I ever use NumPy arrays instead of tf.data?

Yes, for small data that already fits in memory. Passing an array to fit is simpler and fast enough. Switch to tf.data when loading, decoding or augmentation is on the critical path.

Training, saving and serving Mixed precision and multi-device training

Last refreshed 2026-09-18.

The tf.data pipeline in depth

Building a pipeline

Generators, ragged data and lookup tables

Finding the real bottleneck

FAQ

Related