TensorFlow Data Pipelines

You have seen basic tf.data usage in the beginner tutorial. At the intermediate level, the focus shifts to performance building pipelines that keep your GPU saturated, handle large datasets that cannot fit in memory, and apply augmentation efficiently without slowing training.

Introduction to tf.data

tf.data.A dataset represents a lazy sequence of elements. Operations like map(), filter(), and batch() are chained as transformations and only execute when the model requests the next batch. This lazy evaluation is what makes tf.data scalable to arbitrarily large datasets — you never need to hold more than one batch in memory at a time.

Creating Datasets

tf.data supports multiple data sources. The most common are shown below:

import tensorflow as tf

# From NumPy arrays (in-memory)

ds_numpy = tf.data.Dataset.from_tensor_slices((x_train, y_train))

# From image files on disk

ds_images = tf.data.Dataset.list_files('data/images/*.jpg')

# From TFRecord files (most efficient for large datasets)

ds_tfrecord = tf.data.TFRecordDataset('data/train.tfrecord')

# From a Python generator (useful for complex data loading logic)

def generator():

for i in range(1000):

yield (tf.random.normal([28, 28]), i % 10)

ds_gen = tf.data.Dataset.from_generator(

generator,

output_signature=(

tf.TensorSpec(shape=(28,28), dtype=tf.float32),

tf.TensorSpec(shape=(),   dtype=tf.int32)

)

)

Batching and Shuffling Data

The standard pipeline order is: map → cache → shuffle → batch → prefetch. The cache() call stores preprocessed data after the first epoch, avoiding redundant work. According to PythonGuides' deep tf.data analysis, batching after shuffling is non-negotiable — shuffle first, then batch:

BATCH_SIZE  = 64

BUFFER_SIZE = 10000

train_ds = (

tf.data.Dataset.from_tensor_slices((x_train, y_train))

.map(preprocess, num_parallel_calls=tf.data.AUTOTUNE)

.cache()                     # cache after preprocessing

.shuffle(BUFFER_SIZE)        # shuffle individual samples

.batch(BATCH_SIZE)           # group into batches

.prefetch(tf.data.AUTOTUNE)  # overlap data prep with training

)

Data Augmentation

At the intermediate level, augmentation goes beyond simple flips and rotations. You can apply MixUp, CutMix, or random erasing. Here is a practical augmentation pipeline for image classification:

augmentation = tf.keras.Sequential([

tf.keras.layers.RandomFlip('horizontal_and_vertical'),

tf.keras.layers.RandomRotation(0.15),

tf.keras.layers.RandomZoom((-0.1, 0.1)),

tf.keras.layers.RandomTranslation(0.1, 0.1),

tf.keras.layers.RandomContrast(0.2),

tf.keras.layers.RandomBrightness(0.2),

], name='augmentation')

def augment(image, label):

image = tf.cast(image, tf.float32) / 255.0

image = augmentation(image, training=True)

return image, label

train_ds = (

raw_train_ds

.map(augment, num_parallel_calls=tf.data.AUTOTUNE)

.shuffle(5000)

.batch(64)

.prefetch(tf.data.AUTOTUNE)

)

Optimizing Input Pipelines

A poorly built pipeline leaves your GPU waiting for data. Here are the most impactful optimisations: