Introduction

Large TensorFlow datasets should usually be streamed from disk instead of loaded fully into memory. The standard pattern is to keep train, validation, and test splits as separate files or shards, then build a dedicated tf.data pipeline for each split. That gives you predictable evaluation behavior and much better memory usage.

Split the Data on Disk First

For large projects, do not read one giant file into Python and then slice it into train, validation, and test arrays. Keep the splits separate on disk.

1data/
2  train-000.tfrecord
3  train-001.tfrecord
4  valid-000.tfrecord
5  test-000.tfrecord

This makes the split reproducible and lets each pipeline behave differently. Training can shuffle aggressively. Validation and test can remain deterministic.

TFRecord Is a Good Default Format

TensorFlow works especially well with TFRecord because it streams efficiently and integrates cleanly with tf.data.

1import tensorflow as tf
2
3def make_example(features, label):
4    record = tf.train.Example(
5        features=tf.train.Features(
6            feature={
7                "x": tf.train.Feature(float_list=tf.train.FloatList(value=features)),
8                "y": tf.train.Feature(int64_list=tf.train.Int64List(value=[label])),
9            }
10        )
11    )
12    return record.SerializeToString()
13
14with tf.io.TFRecordWriter("train-000.tfrecord") as writer:
15    writer.write(make_example([1.0, 2.0, 3.0], 0))
16    writer.write(make_example([4.0, 5.0, 6.0], 1))

In real projects, use multiple shard files instead of one huge file so TensorFlow can interleave reads more effectively.

Build One Parser and Reuse It

Once the files exist, define one parsing function and reuse it across train, validation, and test pipelines.

1import tensorflow as tf
2
3feature_spec = {
4    "x": tf.io.FixedLenFeature([3], tf.float32),
5    "y": tf.io.FixedLenFeature([], tf.int64),
6}
7
8def parse_record(serialized):
9    example = tf.io.parse_single_example(serialized, feature_spec)
10    return example["x"], example["y"]
11
12def make_dataset(pattern, batch_size, training):
13    files = tf.data.Dataset.list_files(pattern, shuffle=training)
14    dataset = files.interleave(
15        lambda path: tf.data.TFRecordDataset(path),
16        cycle_length=tf.data.AUTOTUNE,
17        num_parallel_calls=tf.data.AUTOTUNE,
18    )
19    dataset = dataset.map(parse_record, num_parallel_calls=tf.data.AUTOTUNE)
20
21    if training:
22        dataset = dataset.shuffle(10000)
23
24    return dataset.batch(batch_size).prefetch(tf.data.AUTOTUNE)
25
26train_ds = make_dataset("data/train-*.tfrecord", batch_size=128, training=True)
27valid_ds = make_dataset("data/valid-*.tfrecord", batch_size=128, training=False)
28test_ds = make_dataset("data/test-*.tfrecord", batch_size=128, training=False)

This pattern scales much better than wrapping large NumPy arrays with from_tensor_slices.

Treat the Three Splits Differently

The pipelines should not be identical.

• training data usually needs shuffling and sometimes augmentation
• validation data should be stable and repeatable
• test data should be deterministic and untouched by training randomness

That is why separate pipelines are worth the effort. Reusing one “do everything” pipeline across all splits often introduces evaluation mistakes.

Non-TFRecord Sources Still Fit the Same Pattern

If the raw data is CSV, text, or images, the same idea still applies: stream from files and transform lazily.

1import tensorflow as tf
2
3train_ds = tf.data.experimental.make_csv_dataset(
4    "data/train.csv",
5    batch_size=256,
6    label_name="target",
7    num_epochs=1,
8    shuffle=True,
9)

TFRecord is not mandatory, but for large TensorFlow-native training workloads it is often the most scalable choice.

Feed the Pipelines Directly Into Training

tf.data.Dataset objects plug directly into Keras.

1model = tf.keras.Sequential([
2    tf.keras.layers.Input(shape=(3,)),
3    tf.keras.layers.Dense(16, activation="relu"),
4    tf.keras.layers.Dense(2),
5])
6
7model.compile(
8    optimizer="adam",
9    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
10    metrics=["accuracy"],
11)
12
13model.fit(train_ds, validation_data=valid_ds, epochs=5)
14model.evaluate(test_ds)

That keeps the entire workflow streaming and avoids unnecessary in-memory copies.

Common Pitfalls

• Loading huge datasets into NumPy first and then running out of RAM.
• Using one shuffled pipeline for validation and test instead of keeping evaluation deterministic.
• Writing too few shard files and limiting I/O throughput.
• Mixing train, validation, and test logic into one pipeline that is hard to reason about.
• Forgetting that preprocessing order affects performance just as much as the model does.

Summary

• Keep large train, validation, and test splits separate on disk.
• Use tf.data to stream data instead of loading everything into memory.
• TFRecord is a strong default for large TensorFlow workloads.
• Give training and evaluation different pipeline behavior on purpose.
• Feed tf.data.Dataset objects directly into model.fit and model.evaluate for scalable training.