Introduction

When your dataset no longer fits in memory, the way you feed data into your model becomes the bottleneck, not the model itself. TensorFlow provides a complete data pipeline framework built around tf.data.Dataset that lets you stream, transform, and prefetch data so the GPU never sits idle waiting for the next batch. This article covers five techniques for efficiently handling large datasets: the tf.data API, TFRecord format, parallel I/O with interleave and prefetch, CSV ingestion, and distributed training with MirroredStrategy.

tf.data.Dataset Pipeline

The tf.data.Dataset API is the foundation of all large-scale data handling in TensorFlow. Instead of loading everything into a NumPy array, you build a lazy pipeline that reads data on demand.

1import tensorflow as tf
2
3# Create a dataset from a directory of images
4dataset = tf.data.Dataset.list_files("data/train/*.jpg", shuffle=True)
5
6def load_and_preprocess(file_path):
7    raw = tf.io.read_file(file_path)
8    image = tf.image.decode_jpeg(raw, channels=3)
9    image = tf.image.resize(image, [224, 224])
10    image = image / 255.0  # Normalize to [0, 1]
11    label = tf.strings.split(file_path, os.sep)[-2]
12    return image, label
13
14dataset = (
15    dataset
16    .map(load_and_preprocess, num_parallel_calls=tf.data.AUTOTUNE)
17    .batch(32)
18    .prefetch(tf.data.AUTOTUNE)
19)

The key idea is that map, batch, and prefetch are chained into a pipeline that overlaps data loading with model training. The AUTOTUNE flag lets TensorFlow dynamically tune the number of parallel threads and the prefetch buffer size based on available resources.

TFRecord Format

TFRecord is TensorFlow's native binary format for storing serialized protocol buffers. It is designed for sequential reads, which makes it significantly faster than reading thousands of individual files from disk.

1# Writing TFRecords
2def serialize_example(image_bytes, label):
3    feature = {
4        'image': tf.train.Feature(
5            bytes_list=tf.train.BytesList(value=[image_bytes])
6        ),
7        'label': tf.train.Feature(
8            int64_list=tf.train.Int64List(value=[label])
9        ),
10    }
11    proto = tf.train.Example(
12        features=tf.train.Features(feature=feature)
13    )
14    return proto.SerializeToString()
15
16with tf.io.TFRecordWriter('train.tfrecord') as writer:
17    for image_bytes, label in data_generator():
18        writer.write(serialize_example(image_bytes, label))
19
20# Reading TFRecords
21def parse_example(serialized):
22    features = tf.io.parse_single_example(serialized, {
23        'image': tf.io.FixedLenFeature([], tf.string),
24        'label': tf.io.FixedLenFeature([], tf.int64),
25    })
26    image = tf.io.decode_jpeg(features['image'], channels=3)
27    image = tf.image.resize(image, [224, 224]) / 255.0
28    return image, features['label']
29
30dataset = (
31    tf.data.TFRecordDataset('train.tfrecord')
32    .map(parse_example, num_parallel_calls=tf.data.AUTOTUNE)
33    .shuffle(1000)
34    .batch(32)
35    .prefetch(tf.data.AUTOTUNE)
36)

For very large datasets, split the data across multiple TFRecord files (called sharding). This enables parallel reads and makes it easier to distribute data across machines.

Interleave and Prefetch for Parallel I/O

When data is spread across many files, interleave reads from multiple files simultaneously. Combined with prefetch, this keeps the training loop fed with data even when individual file reads are slow.

1file_pattern = "data/train-*.tfrecord"
2files = tf.data.Dataset.list_files(file_pattern)
3
4dataset = files.interleave(
5    lambda f: tf.data.TFRecordDataset(f),
6    cycle_length=4,              # Read from 4 files at once
7    num_parallel_calls=tf.data.AUTOTUNE,
8    deterministic=False          # Allow out-of-order for speed
9)
10
11dataset = (
12    dataset
13    .map(parse_example, num_parallel_calls=tf.data.AUTOTUNE)
14    .shuffle(buffer_size=10000)
15    .batch(64)
16    .prefetch(tf.data.AUTOTUNE)
17)

The cycle_length parameter controls how many files are open at the same time. Setting deterministic=False allows TensorFlow to yield whichever record is ready first, which improves throughput when file read times vary. The prefetch at the end ensures that the next batch is already in memory while the current batch is being processed by the GPU.

CsvDataset for Tabular Data

For tabular data stored in CSV files, tf.data.experimental.CsvDataset reads rows directly into tensors without loading the entire file into a pandas DataFrame first.

1# Define column types (one per column)
2column_defaults = [tf.float32, tf.float32, tf.float32, tf.int32]
3
4dataset = tf.data.experimental.CsvDataset(
5    filenames=["data/train.csv"],
6    record_defaults=column_defaults,
7    header=True,
8    select_cols=[0, 1, 2, 3]  # Read only specific columns
9)
10
11def pack_features(col1, col2, col3, label):
12    features = tf.stack([col1, col2, col3])
13    return features, label
14
15dataset = (
16    dataset
17    .map(pack_features)
18    .shuffle(5000)
19    .batch(128)
20    .prefetch(tf.data.AUTOTUNE)
21)

You can also pass a list of filenames to read from multiple CSV files. This is useful when data arrives in daily or hourly partitions. For more complex CSV processing (missing values, string encoding), consider using tf.data.experimental.make_csv_dataset, which handles column naming and batching automatically.

Distributed Training with MirroredStrategy

When a single GPU cannot process data fast enough, tf.distribute.MirroredStrategy replicates your model across multiple GPUs on the same machine. Each GPU processes a slice of the batch in parallel, and gradients are synchronized automatically.

1strategy = tf.distribute.MirroredStrategy()
2print(f"Number of devices: {strategy.num_replicas_in_sync}")
3
4# Scale the batch size by the number of GPUs
5GLOBAL_BATCH_SIZE = 64 * strategy.num_replicas_in_sync
6
7dataset = (
8    tf.data.TFRecordDataset(tf.io.gfile.glob("data/train-*.tfrecord"))
9    .map(parse_example, num_parallel_calls=tf.data.AUTOTUNE)
10    .shuffle(10000)
11    .batch(GLOBAL_BATCH_SIZE)
12    .prefetch(tf.data.AUTOTUNE)
13)
14
15with strategy.scope():
16    model = tf.keras.Sequential([
17        tf.keras.layers.Conv2D(32, 3, activation='relu',
18                               input_shape=(224, 224, 3)),
19        tf.keras.layers.GlobalAveragePooling2D(),
20        tf.keras.layers.Dense(10, activation='softmax'),
21    ])
22    model.compile(
23        optimizer='adam',
24        loss='sparse_categorical_crossentropy',
25        metrics=['accuracy']
26    )
27
28model.fit(dataset, epochs=10)

The critical detail is that everything inside strategy.scope() is replicated. The dataset is automatically sharded across GPUs. Scale the global batch size proportionally so each GPU still processes the same per-replica batch size. For multi-machine training, switch to tf.distribute.MultiWorkerMirroredStrategy.

Common Pitfalls

• Loading the entire dataset into memory with NumPy before creating a Dataset: This defeats the purpose of tf.data. Use list_files, TFRecordDataset, or CsvDataset to stream data from disk.
• Forgetting to call prefetch at the end of the pipeline: Without prefetch, the GPU waits for data loading to finish after each batch. A single .prefetch(tf.data.AUTOTUNE) at the end of your pipeline is the simplest performance win.
• Setting shuffle buffer_size too small: A buffer of 100 on a dataset of 1 million records produces nearly sequential batches, which can hurt model convergence. Set the buffer to at least several thousand, or use pre-shuffled TFRecord shards.
• Not scaling the batch size with MirroredStrategy: If you use the same global batch size across 4 GPUs, each GPU only processes one quarter of a batch, wasting compute. Multiply the base batch size by the number of replicas.
• Writing a single massive TFRecord file: One large file cannot be read in parallel. Shard your data into files of 100 to 200 MB each so interleave can read from multiple files simultaneously.

Summary

• tf.data.Dataset is the core abstraction for streaming data through a lazy, chainable pipeline of map, batch, and prefetch operations.
• TFRecord is TensorFlow's optimized binary format for sequential reads; shard large datasets into multiple files for parallel I/O.
• interleave + prefetch overlap file reading with GPU training, eliminating I/O bottlenecks in multi-file pipelines.
• CsvDataset reads tabular data directly into tensors without loading entire files into memory.
• MirroredStrategy distributes training across multiple GPUs with automatic gradient synchronization; always scale the global batch size by the number of replicas.