Loading & Preprocessing Data with tf.data

What you’ll learn

what a tf.data.Datasettf.data.Dataset is, and how transformation methods chain together
how to shuffle, repeat, batch, and preprocess data on the fly with map()map()
why reading from many files and interleaving them beats reading one file at a time
how to store large datasets efficiently using the TFRecord format
why prefetch()prefetch() is the one line that keeps your GPU from sitting idle

Why not just load a NumPy array?

Every dataset so far in this module fit comfortably in RAM. Real datasets — millions of images, gigabytes of logs — don’t. The Data API (tf.datatf.data) solves this by letting you describe a pipeline instead of a fixed array: “read from these files, skip the header, parse each line, shuffle, batch, and get the next batch ready while the GPU is still busy with the current one.” TensorFlow handles the multithreading, queuing, and memory management for you.

The core idea: chaining transformations

A tf.data.Datasettf.data.Dataset represents a sequence of items. Every transformation method (.map().map(), .batch().batch(), .shuffle().shuffle(), .filter().filter(), .repeat().repeat(), …) returns a new dataset — it never modifies the original — so you chain them together:

Chaining dataset transformations

import tensorflow as tf
 
X = tf.range(10)
dataset = tf.data.Dataset.from_tensor_slices(X)
 
# repeat the 10 items 3 times, then group them into batches of 7
dataset = dataset.repeat(3).batch(7)
 
for item in dataset:
    print(item.numpy())
 
# [0 1 2 3 4 5 6]
# [7 8 9 0 1 2 3]
# [4 5 6 7 8 9 0]
# [1 2 3 4 5 6 7]
# [8 9]

Chaining dataset transformations

import tensorflow as tf
 
X = tf.range(10)
dataset = tf.data.Dataset.from_tensor_slices(X)
 
# repeat the 10 items 3 times, then group them into batches of 7
dataset = dataset.repeat(3).batch(7)
 
for item in dataset:
    print(item.numpy())
 
# [0 1 2 3 4 5 6]
# [7 8 9 0 1 2 3]
# [4 5 6 7 8 9 0]
# [1 2 3 4 5 6 7]
# [8 9]

repeat()repeat() with no argument repeats forever — the code that iterates decides when to stop. batch()batch() outputs a smaller final batch unless you pass drop_remainder=Truedrop_remainder=True.

Use .map().map() to apply preprocessing to every item (pass num_parallel_callsnum_parallel_calls to spawn multiple threads for CPU-heavy work like resizing images), and .filter().filter() to drop items you don’t want:

map() and filter()

dataset = tf.data.Dataset.range(10)
dataset = dataset.map(lambda x: x * 2, num_parallel_calls=tf.data.AUTOTUNE)
dataset = dataset.filter(lambda x: x < 10)
 
print(list(dataset.as_numpy_iterator()))
# [0, 2, 4, 6, 8]

map() and filter()

dataset = tf.data.Dataset.range(10)
dataset = dataset.map(lambda x: x * 2, num_parallel_calls=tf.data.AUTOTUNE)
dataset = dataset.filter(lambda x: x < 10)
 
print(list(dataset.as_numpy_iterator()))
# [0, 2, 4, 6, 8]

Shuffling

Gradient Descent works best when training instances are independent and identically distributed, so shuffle()shuffle() fills a buffer with items from the source dataset and pulls out a random one every time it’s asked for an item, refilling from the source as it goes:

Shuffling with a buffer

dataset = tf.data.Dataset.range(10).repeat(3)
dataset = dataset.shuffle(buffer_size=5, seed=42).batch(7)
 
for item in dataset:
    print(item.numpy())

Shuffling with a buffer

dataset = tf.data.Dataset.range(10).repeat(3)
dataset = dataset.shuffle(buffer_size=5, seed=42).batch(7)
 
for item in dataset:
    print(item.numpy())

The buffer must be large enough to shuffle effectively — but it doesn’t need to exceed your dataset’s size. For datasets too big to shuffle well in memory, split the data across many files, read several of them at once with interleave()interleave(), and layer a shuffle()shuffle() buffer on top of that.

Reading and interleaving many files

Real pipelines usually read from many files at once, cycling through them so records from different files get interleaved instead of arriving in file order:

Interleaving multiple CSV files

train_filepaths = ["data/housing_00.csv", "data/housing_01.csv", "data/housing_02.csv"]
 
filepath_dataset = tf.data.Dataset.list_files(train_filepaths, seed=42)
 
n_readers = 5
dataset = filepath_dataset.interleave(
    lambda filepath: tf.data.TextLineDataset(filepath).skip(1),  # skip header row
    cycle_length=n_readers,
    num_parallel_calls=tf.data.AUTOTUNE,
)

Interleaving multiple CSV files

train_filepaths = ["data/housing_00.csv", "data/housing_01.csv", "data/housing_02.csv"]
 
filepath_dataset = tf.data.Dataset.list_files(train_filepaths, seed=42)
 
n_readers = 5
dataset = filepath_dataset.interleave(
    lambda filepath: tf.data.TextLineDataset(filepath).skip(1),  # skip header row
    cycle_length=n_readers,
    num_parallel_calls=tf.data.AUTOTUNE,
)

Putting it together: a reusable CSV pipeline

A full csv_reader_dataset() pipeline

import tensorflow as tf
 
n_inputs = 8  # number of feature columns
 
def preprocess(line, X_mean, X_std):
    defs = [0.0] * n_inputs + [tf.constant([], dtype=tf.float32)]
    fields = tf.io.decode_csv(line, record_defaults=defs)
    x = tf.stack(fields[:-1])
    y = tf.stack(fields[-1:])
    return (x - X_mean) / X_std, y
 
def csv_reader_dataset(
    filepaths, X_mean, X_std, repeat=1, n_readers=5,
    shuffle_buffer_size=10_000, n_parse_threads=5, batch_size=32,
):
    dataset = tf.data.Dataset.list_files(filepaths)
    dataset = dataset.interleave(
        lambda fp: tf.data.TextLineDataset(fp).skip(1),
        cycle_length=n_readers,
        num_parallel_calls=tf.data.AUTOTUNE,
    )
    dataset = dataset.map(
        lambda line: preprocess(line, X_mean, X_std),
        num_parallel_calls=n_parse_threads,
    )
    dataset = dataset.shuffle(shuffle_buffer_size).repeat(repeat)
    return dataset.batch(batch_size).prefetch(1)
 
# train_set = csv_reader_dataset(train_filepaths, X_mean, X_std)
# model.fit(train_set, epochs=10, validation_data=valid_set)

A full csv_reader_dataset() pipeline

import tensorflow as tf
 
n_inputs = 8  # number of feature columns
 
def preprocess(line, X_mean, X_std):
    defs = [0.0] * n_inputs + [tf.constant([], dtype=tf.float32)]
    fields = tf.io.decode_csv(line, record_defaults=defs)
    x = tf.stack(fields[:-1])
    y = tf.stack(fields[-1:])
    return (x - X_mean) / X_std, y
 
def csv_reader_dataset(
    filepaths, X_mean, X_std, repeat=1, n_readers=5,
    shuffle_buffer_size=10_000, n_parse_threads=5, batch_size=32,
):
    dataset = tf.data.Dataset.list_files(filepaths)
    dataset = dataset.interleave(
        lambda fp: tf.data.TextLineDataset(fp).skip(1),
        cycle_length=n_readers,
        num_parallel_calls=tf.data.AUTOTUNE,
    )
    dataset = dataset.map(
        lambda line: preprocess(line, X_mean, X_std),
        num_parallel_calls=n_parse_threads,
    )
    dataset = dataset.shuffle(shuffle_buffer_size).repeat(repeat)
    return dataset.batch(batch_size).prefetch(1)
 
# train_set = csv_reader_dataset(train_filepaths, X_mean, X_std)
# model.fit(train_set, epochs=10, validation_data=valid_set)

Notice the pipeline order: read → preprocess → shuffle → repeat → batch → prefetch. If the whole preprocessed dataset fits in RAM, insert .cache().cache() right after preprocessing (before shuffling) so each instance is only parsed once per run, not once per epoch.

The TFRecord format

CSV is simple but not very efficient, and it doesn’t handle complex data like images well. TFRecord is TensorFlow’s own binary format: a flat sequence of records, usually serialized ExampleExample protocol buffers.

Writing and reading a TFRecord file

import tensorflow as tf
 
# Write a couple of raw records
with tf.io.TFRecordWriter("my_data.tfrecord") as f:
    f.write(b"This is the first record")
    f.write(b"And this is the second record")
 
# Read them back
dataset = tf.data.TFRecordDataset(["my_data.tfrecord"])
for item in dataset:
    print(item.numpy())

Writing and reading a TFRecord file

import tensorflow as tf
 
# Write a couple of raw records
with tf.io.TFRecordWriter("my_data.tfrecord") as f:
    f.write(b"This is the first record")
    f.write(b"And this is the second record")
 
# Read them back
dataset = tf.data.TFRecordDataset(["my_data.tfrecord"])
for item in dataset:
    print(item.numpy())

In practice you serialize structured data using the ExampleExample protobuf, one per instance, with named floatfloat, int64int64, or bytesbytes features:

Serializing an Example protobuf

from tensorflow.train import BytesList, FloatList, Int64List
from tensorflow.train import Feature, Features, Example
import tensorflow as tf
 
person_example = Example(
    features=Features(
        feature={
            "name": Feature(bytes_list=BytesList(value=[b"Alice"])),
            "id": Feature(int64_list=Int64List(value=[123])),
            "emails": Feature(bytes_list=BytesList(value=[b"a@b.com", b"c@d.com"])),
        }
    )
)
 
with tf.io.TFRecordWriter("my_contacts.tfrecord") as f:
    f.write(person_example.SerializeToString())
 
# Parsing it back requires a feature description
feature_description = {
    "name": tf.io.FixedLenFeature([], tf.string, default_value=""),
    "id": tf.io.FixedLenFeature([], tf.int64, default_value=0),
    "emails": tf.io.VarLenFeature(tf.string),
}
 
for serialized in tf.data.TFRecordDataset(["my_contacts.tfrecord"]):
    parsed = tf.io.parse_single_example(serialized, feature_description)
    print(parsed["name"].numpy(), tf.sparse.to_dense(parsed["emails"]).numpy())

Serializing an Example protobuf

from tensorflow.train import BytesList, FloatList, Int64List
from tensorflow.train import Feature, Features, Example
import tensorflow as tf
 
person_example = Example(
    features=Features(
        feature={
            "name": Feature(bytes_list=BytesList(value=[b"Alice"])),
            "id": Feature(int64_list=Int64List(value=[123])),
            "emails": Feature(bytes_list=BytesList(value=[b"a@b.com", b"c@d.com"])),
        }
    )
)
 
with tf.io.TFRecordWriter("my_contacts.tfrecord") as f:
    f.write(person_example.SerializeToString())
 
# Parsing it back requires a feature description
feature_description = {
    "name": tf.io.FixedLenFeature([], tf.string, default_value=""),
    "id": tf.io.FixedLenFeature([], tf.int64, default_value=0),
    "emails": tf.io.VarLenFeature(tf.string),
}
 
for serialized in tf.data.TFRecordDataset(["my_contacts.tfrecord"]):
    parsed = tf.io.parse_single_example(serialized, feature_description)
    print(parsed["name"].numpy(), tf.sparse.to_dense(parsed["emails"]).numpy())

Set compression_type="GZIP"compression_type="GZIP" on both TFRecordOptionsTFRecordOptions (when writing) and TFRecordDatasetTFRecordDataset (when reading) if the files travel over a slow network.

Prefetching: the one line that matters most

diagram Diagram mermaid

  flowchart LR
  subgraph "Without prefetch"
    A1["CPU: load batch 1"] --> A2["GPU: train batch 1"] --> A3["CPU: load batch 2"] --> A4["GPU: train batch 2"]
  end

diagram Diagram mermaid

  flowchart LR
  subgraph "With prefetch(1)"
    B1["CPU: load batch 1"] --> B2["GPU: train batch 1"]
    B1 --> C1["CPU: load batch 2 (in parallel)"]
    B2 --> C2["GPU: train batch 2"]
    C1 --> C2
    C1 --> D1["CPU: load batch 3 (in parallel)"]
  end

Calling .prefetch(1).prefetch(1) at the end of a pipeline tells tf.datatf.data to always try to have the next batch ready before it’s needed. While the GPU trains on the current batch, the CPU is already loading and preprocessing the next one in a background thread — so the GPU (usually the most expensive, most limited resource) is almost never idle.

Using a dataset with tf.keras

Once you have a pipeline, pass it straight to fit()fit(), evaluate()evaluate(), or predict()predict() instead of separate XX/yy arrays — tf.kerastf.keras understands tf.data.Datasettf.data.Dataset objects natively:

Training with a tf.data pipeline

# train_set = csv_reader_dataset(train_filepaths, X_mean, X_std)
# valid_set = csv_reader_dataset(valid_filepaths, X_mean, X_std)
#
# model = keras.models.Sequential([...])
# model.compile(loss="mse", optimizer="adam")
# history = model.fit(train_set, epochs=10, validation_data=valid_set)

Training with a tf.data pipeline

# train_set = csv_reader_dataset(train_filepaths, X_mean, X_std)
# valid_set = csv_reader_dataset(valid_filepaths, X_mean, X_std)
#
# model = keras.models.Sequential([...])
# model.compile(loss="mse", optimizer="adam")
# history = model.fit(train_set, epochs=10, validation_data=valid_set)

Letting TensorFlow pick the buffer size

Every num_parallel_callsnum_parallel_calls and prefetch()prefetch() example so far uses a hand-picked number. If you’re not sure what to pick, pass tf.data.AUTOTUNEtf.data.AUTOTUNE instead of a fixed integer — TensorFlow tunes the value dynamically at runtime, based on the resources actually available:

Letting TensorFlow tune prefetch and parallelism automatically

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

Letting TensorFlow tune prefetch and parallelism automatically

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

Visualize it

Watch the CPU and GPU work in lockstep once prefetching kicks in — the CPU is always one batch ahead, so the GPU never has to wait:

sketch Prefetching: CPU stays one batch ahead of the GPU p5.js

A queue fills with batches prepared by the CPU while the GPU consumes one batch at a time; with prefetching the GPU rarely stalls.

Mini-checkpoint

If your training loop is GPU-bound but the GPU utilization graph shows regular dips:

your input pipeline is probably the bottleneck — add num_parallel_callsnum_parallel_calls to interleave()interleave()/map()map(), and make sure you’re calling .prefetch().prefetch().

🧪 Try It Yourself

Exercise 1 – Chain repeat() and batch()

Exercise 2 – Shuffle with a buffer

Exercise 3 – Add prefetch()

Continue to Custom Models and Training Loops (TensorFlow) — once your model is fed by a pipeline like this, the next step is looking at what’s happening inside fit()fit() itself, and how to customize it when you need to.

If this helped you, consider buying me a coffee ☕

Buy me a coffee