tf.data map, shuffle, batch, cache and prefetch Explained

The Story — A Restaurant Kitchen at Rush Hour

A tf.data pipeline is a chain of these stations. Each method returns a new dataset, so you build the pipeline step by step. The order of the steps changes both what you get and how fast you get it.

map, filter and batch

import tensorflow as tf
AUTOTUNE = tf.data.AUTOTUNE

ds = tf.data.Dataset.range(10)
ds = ds.map(lambda x: x * x, num_parallel_calls=AUTOTUNE)   # square each value
ds = ds.filter(lambda x: x % 2 == 0)                        # keep even squares
print("after map+filter:", [int(v) for v in ds])

print("batch(3)        :", [b.numpy().tolist() for b in tf.data.Dataset.range(8).batch(3)])
print("drop_remainder  :", [b.numpy().tolist() for b in tf.data.Dataset.range(8).batch(3, drop_remainder=True)])

shuffle — The Buffer Trick

A dataset may be too big to load all at once, so shuffle(n) keeps a buffer of n elements. Each time you ask for an element, it picks a random one from the buffer and refills the gap with the next element from the input. A small buffer only mixes nearby elements. For a full shuffle, the buffer must be at least as big as the dataset.

base = tf.data.Dataset.range(12)
print("buffer 1  :", [int(v) for v in base.shuffle(1, seed=0)])
print("buffer 3  :", [int(v) for v in base.shuffle(3, seed=0)])
print("buffer 12 :", [int(v) for v in base.shuffle(12, seed=0)])

Order Matters: shuffle Before batch

print("shuffle then batch:", [b.numpy().tolist() for b in tf.data.Dataset.range(9).shuffle(9, seed=1).batch(3)])
print("batch then shuffle:", [b.numpy().tolist() for b in tf.data.Dataset.range(9).batch(3).shuffle(3, seed=1)])

Batching first keeps the same neighbours together forever; only the batch order changes. Shuffle before you batch.

Speed: parallel map, cache and prefetch

We fake a slow disk read (20 ms per element) and a fast training step (5 ms). The data is the bottleneck. Then we time two epochs with each pipeline. TensorFlow quietly adds some speed tricks by default, so plain() switches them off to give a fair starting point.

import time

def slow_read(i):
    time.sleep(0.02)                    # pretend to read a file
    return i

def load(i):
    return tf.py_function(slow_read, [i], tf.int64)

def plain(ds):                          # turn off automatic tricks for a fair test
    opts = tf.data.Options()
    opts.autotune.enabled = False
    opts.experimental_optimization.inject_prefetch = False
    return ds.with_options(opts)

def time_two_epochs(ds):
    t0 = time.perf_counter()
    for _ in range(2):
        for _ in ds:
            time.sleep(0.005)           # pretend to train on it
    return time.perf_counter() - t0

base = tf.data.Dataset.range(30)
pipelines = {
    "map                     ": base.map(load),
    "map + prefetch          ": base.map(load).prefetch(1),
    "parallel map + prefetch ": base.map(load, num_parallel_calls=4).prefetch(1),
    "parallel + cache + pref.": base.map(load, num_parallel_calls=4).cache().prefetch(1),
}
for name, ds in pipelines.items():
    print(f"{name}: {time_two_epochs(plain(ds)):.2f} s")

Each stage helps: prefetch hides the training time behind the reading, parallel map reads 4 files at once, and cache skips the reading completely in epoch 2.

Try It — Pipeline Timeline Builder

The Recommended Order

ds = (tf.data.Dataset.from_tensor_slices((paths, labels))
        .map(load_and_decode, num_parallel_calls=AUTOTUNE)   # 1. expensive, same every epoch
        .cache()                                             # 2. remember the decoded data
        .shuffle(10_000)                                     # 3. new order every epoch
        .map(random_augment, num_parallel_calls=AUTOTUNE)    # 4. random -> must be AFTER cache
        .batch(32)                                           # 5. group
        .prefetch(AUTOTUNE))                                 # 6. always last

Golden Rules