TensorFlow Data Pipeline

A data pipeline is the system that reads, preprocesses, and feeds data to your model during training. A poorly designed pipeline can make the GPU sit idle 80% of the time waiting for data — wasting expensive compute. The tf.data API solves this with a fast, memory-efficient pipeline that reads data in parallel, preprocesses it on the CPU while the GPU trains, and eliminates the gap between batches. Mastering tf.data separates amateur TensorFlow projects from production-grade systems.

The Factory Production Line Analogy

Without a proper pipeline, the factory (GPU) processes one box of parts (batch), then stops and waits while the warehouse staff (CPU) slowly unpacks the next box. With tf.data, the warehouse staff prepare ten boxes in advance while the factory processes the current one. The factory never stops. Throughput doubles or triples with no hardware changes.

Creating a Dataset From Arrays

import tensorflow as tf
import numpy as np

# Create sample data
x = np.random.random((1000, 20)).astype('float32')
y = np.random.randint(0, 3, size=(1000,))

# Wrap in a tf.data.Dataset
dataset = tf.data.Dataset.from_tensor_slices((x, y))
print(dataset)
# 

# Iterate through the dataset
for features, label in dataset.take(3):
    print(features.shape, label.numpy())

The Four Essential Pipeline Operations

1. shuffle() — Randomize Order

Shuffling prevents the model from learning the order of examples rather than the patterns within them. A shuffle buffer holds a fixed number of samples in memory and randomly draws from them.

# buffer_size should be >= dataset size for full shuffle
# Use a smaller number if memory is limited
dataset = dataset.shuffle(buffer_size=1000, seed=42)

2. batch() — Group Into Batches

Batching groups individual examples into tensors that can be processed in parallel on the GPU. Larger batches use memory more efficiently but produce noisier gradient estimates.

dataset = dataset.batch(32)
# Each element is now (batch_of_32_features, batch_of_32_labels)

3. prefetch() — Overlap CPU and GPU Work

Prefetching prepares the next batch on the CPU while the GPU processes the current one. AUTOTUNE tells TensorFlow to tune the prefetch buffer size automatically based on available memory and throughput.

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

4. map() — Transform Each Example

The map() function applies a preprocessing function to each element. Use it to normalize values, parse records, apply augmentation, or convert data types.

def normalize(features, label):
    features = tf.cast(features, tf.float32) / 255.0
    return features, label

dataset = dataset.map(normalize, num_parallel_calls=tf.data.AUTOTUNE)

The Standard Training Pipeline Pattern

def make_dataset(x, y, batch_size=32, shuffle=True, augment=False):
    dataset = tf.data.Dataset.from_tensor_slices((x, y))

    if shuffle:
        dataset = dataset.shuffle(buffer_size=len(x))

    if augment:
        dataset = dataset.map(augment_fn, num_parallel_calls=tf.data.AUTOTUNE)

    dataset = dataset.map(preprocess_fn, num_parallel_calls=tf.data.AUTOTUNE)
    dataset = dataset.batch(batch_size)
    dataset = dataset.prefetch(tf.data.AUTOTUNE)
    return dataset

train_ds = make_dataset(x_train, y_train, shuffle=True, augment=True)
val_ds   = make_dataset(x_val, y_val, shuffle=False, augment=False)
test_ds  = make_dataset(x_test, y_test, shuffle=False, augment=False)

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

Pipeline Diagram

Raw Data (disk or memory)
         │
         ▼
[from_tensor_slices() or from_generator() or list_files()]
"Create dataset — one element per sample"
         │
         ▼
[shuffle(buffer_size=N)]
"Randomize order — fills buffer and draws randomly"
         │
         ▼
[map(preprocess_fn, num_parallel_calls=AUTOTUNE)]
"Apply transforms in parallel on multiple CPU cores"
         │
         ▼
[batch(32)]
"Group 32 samples into one tensor"
         │
         ▼
[prefetch(AUTOTUNE)]
"Pre-prepare next batch while GPU processes current batch"
         │
         ▼
[GPU Training]

Loading Data From Files

From a List of File Paths

import pathlib

data_dir = pathlib.Path('images/')
file_list = tf.data.Dataset.list_files(str(data_dir/'*/*.jpg'), shuffle=True)

def load_image(path):
    image = tf.io.read_file(path)
    image = tf.image.decode_jpeg(image, channels=3)
    image = tf.image.resize(image, [224, 224])
    image = tf.cast(image, tf.float32) / 255.0
    return image

image_dataset = file_list.map(load_image, num_parallel_calls=tf.data.AUTOTUNE)

From a CSV File

dataset = tf.data.experimental.make_csv_dataset(
    'data.csv',
    batch_size=32,
    label_name='target_column',
    num_epochs=1
)

Caching: Speed Up Repeated Preprocessing

After the first epoch, tf.data can cache the processed data in memory (or on disk). Subsequent epochs skip the preprocessing step entirely, which dramatically speeds up training when preprocessing is expensive.

dataset = (
    tf.data.Dataset.from_tensor_slices((x, y))
    .map(expensive_preprocess, num_parallel_calls=tf.data.AUTOTUNE)
    .cache()          # Cache after preprocessing — only done once
    .shuffle(1000)    # Shuffle after cache so shuffle works each epoch
    .batch(32)
    .prefetch(tf.data.AUTOTUNE)
)

Note: always put cache() before shuffle() and batch(). If you cache after batching, TensorFlow caches the same batches every epoch, making shuffle ineffective.

Measuring Pipeline Throughput

import time

# Benchmark how fast the pipeline delivers batches
dataset = make_dataset(x_train, y_train)
start = time.time()
count = 0
for batch in dataset:
    count += 1
elapsed = time.time() - start

print(f"Delivered {count} batches in {elapsed:.2f}s")
print(f"Throughput: {count/elapsed:.1f} batches/sec")

Performance Comparison

Configuration                   Training speed (relative)
──────────────────────────────────────────────────────────
NumPy arrays fed directly            1.0×  (baseline)
tf.data, no prefetch, no cache       0.8×  (overhead from setup)
tf.data + prefetch                   1.8×
tf.data + cache + prefetch           2.5×
tf.data + parallel map + prefetch    3.2×
tf.data + cache + parallel + prefetch 4.0×  (ideal)
──────────────────────────────────────────────────────────

A well-designed tf.data pipeline feeds the GPU without interruption, turning a slow training run into a fast one without upgrading hardware. The next topic covers loading CSV data specifically — a common real-world task where tabular data from files needs to be converted into tensors for training.