TensorFlow Data Pipelines | TensorFlow, Interactively

Every playground so far has fed a model from a small array already sitting in memory. Real datasets don't fit in memory, need per-example preprocessing (decoding an image, tokenizing text), and — if you're not careful — can leave an expensive GPU sitting idle while the CPU prepares the next batch. tf.data.Dataset is the answer: instead of loading and transforming everything up front, you describe a pipeline, and it pulls data through that pipeline lazily, one batch at a time, as training asks for it.

A typical pipeline chains a handful of operations, and the order they're chained in matters:

.shuffle(buffer_size) — fills a buffer of buffer_size examples and pulls randomly from it, refilling as it goes. It's not a full shuffle of the whole dataset (that would mean loading everything into memory first) — it's a good-enough approximation that trades randomness quality for staying lazy. Bigger buffer, better shuffle, more memory.
.map(preprocess_fn) — applies a transformation to every example: decode an image, resize it, tokenize a string. This is usually where the real CPU cost of a pipeline lives.
.batch(batch_size) — groups consecutive examples into a batch, the unit model.fit actually trains on.
.cache() — stores each example's output the first time it's computed, so every subsequent epoch skips recomputing it. Only worth it when preprocessing is expensive and deterministic — cache a random augmentation and you've cached away the randomness too.
.prefetch(tf.data.AUTOTUNE) — the one that actually keeps the GPU fed: it overlaps preparing the next batch with the model training on the current one, instead of making training wait on preprocessing every single step.

🔍 Deep dive: Why prefetch is the one that actually matters for speed

Without prefetching, a training loop looks like: preprocess batch → train on it → preprocess the next batch → train on it → ... — strictly serial, so total time is roughly (preprocess_time + train_time) × num_batches. With prefetching, the next batch's preprocessing starts while the current one is training — as long as there's spare CPU capacity to do it, total time drops toward max(preprocess_time, train_time) × num_batches instead. The playground below measures this directly: a real tf.data pipeline, a real (simulated-cost) preprocessing step, and real wall-clock milliseconds, with prefetching toggled on and off.

The playground runs an actual tf.data.Dataset pipeline in your browser (tfjs's own tf.data module, the same API family as Python's) — .array(), .shuffle(), .mapAsync(), .prefetch() are all real calls, not a mock of them. The one gap: tfjs's tf.data doesn't have a .cache() method (only the Python API does), so caching here is a small hand-rolled stand-in — a shared lookup table checked inside the preprocessing step — that behaves identically.

Production note

tf.data.AUTOTUNE (in the buffer-size position of .prefetch()) tells TensorFlow to pick the buffer size dynamically at runtime based on available resources, instead of you guessing a fixed number. It's almost always the right default.

Beginner tip

If two settings in the playground produce nearly the same time, that's not a bug — prefetching can only hide preprocessing cost up to however long the training step itself takes. Once preprocessing is faster than training, there's nothing left to hide.