Introduction

In TensorFlow 2.x, the standard way to feed data into a model is the tf.data API. If you come from PyTorch, you may look for a separate “DataLoader” object, but in TensorFlow the tf.data.Dataset pipeline usually plays that role by handling loading, transformation, batching, and prefetching.

The Role of tf.data.Dataset

A Dataset is an iterable sequence of elements. Each element can be a single tensor, a tuple such as features and labels, or a nested structure. The usual pipeline is:

• create a dataset
• transform it with operations like map or filter
• randomize it with shuffle
• group records with batch
• overlap input work with model execution using prefetch

For small in-memory arrays, from_tensor_slices is the simplest starting point:

1import numpy as np
2import tensorflow as tf
3
4x = np.array([[1.0], [2.0], [3.0], [4.0]], dtype=np.float32)
5y = np.array([2.0, 4.0, 6.0, 8.0], dtype=np.float32)
6
7dataset = tf.data.Dataset.from_tensor_slices((x, y))
8
9for features, label in dataset.take(2):
10    print(features.numpy(), label.numpy())

Each item yielded by the dataset contains one input sample and one target value.

Building a Training Pipeline

The power of tf.data comes from chaining transformations. A realistic training pipeline often looks like this:

1import numpy as np
2import tensorflow as tf
3
4x = np.arange(20, dtype=np.float32).reshape(-1, 1)
5y = 3 * x.squeeze() + 1
6
7dataset = (
8    tf.data.Dataset.from_tensor_slices((x, y))
9    .shuffle(buffer_size=len(x))
10    .batch(4)
11    .prefetch(tf.data.AUTOTUNE)
12)
13
14for batch_x, batch_y in dataset.take(1):
15    print(batch_x.shape)
16    print(batch_y.shape)

This does three important things:

• 'shuffle prevents the model from always seeing the same order'
• 'batch groups examples into mini-batches'
• 'prefetch prepares the next batch while the current one is being consumed'

For many applications, this is the TensorFlow equivalent of “dataset plus dataloader.”

Transforming Data with map

The map step is where you parse files, normalize values, tokenize text, or apply image preprocessing.

1import tensorflow as tf
2
3dataset = tf.data.Dataset.from_tensor_slices([1.0, 2.0, 3.0, 4.0])
4
5dataset = dataset.map(lambda value: value / 4.0)
6
7for value in dataset:
8    print(value.numpy())

In real projects, the mapping function is often more complex:

1import tensorflow as tf
2
3features = tf.data.Dataset.from_tensor_slices(["1.0,2.0", "3.0,4.0"])
4
5def parse_csv_line(line):
6    values = tf.strings.to_number(tf.strings.split(line, ","), tf.float32)
7    return values
8
9parsed = features.map(parse_csv_line, num_parallel_calls=tf.data.AUTOTUNE)
10
11for row in parsed:
12    print(row.numpy())

The num_parallel_calls option lets TensorFlow process several elements in parallel when the transformation is expensive.

Feeding a Keras Model

Keras can consume a Dataset directly. That is one of the main reasons tf.data is so useful.

1import numpy as np
2import tensorflow as tf
3
4x = np.arange(100, dtype=np.float32).reshape(-1, 1)
5y = 2 * x.squeeze() + 5
6
7train_ds = (
8    tf.data.Dataset.from_tensor_slices((x, y))
9    .shuffle(100)
10    .batch(16)
11    .prefetch(tf.data.AUTOTUNE)
12)
13
14model = tf.keras.Sequential([
15    tf.keras.layers.Input(shape=(1,)),
16    tf.keras.layers.Dense(1)
17])
18
19model.compile(optimizer="adam", loss="mse")
20model.fit(train_ds, epochs=3, verbose=0)

Notice that there is no separate loader object. The dataset itself describes how the data should be delivered.

When Data Does Not Fit in Memory

For larger workloads, you usually build the dataset from files or generators instead of NumPy arrays.

Here is a minimal generator-based example:

1import tensorflow as tf
2
3def number_generator():
4    for i in range(5):
5        yield i, i * i
6
7dataset = tf.data.Dataset.from_generator(
8    number_generator,
9    output_signature=(
10        tf.TensorSpec(shape=(), dtype=tf.int32),
11        tf.TensorSpec(shape=(), dtype=tf.int32),
12    ),
13).batch(2)
14
15for batch_x, batch_y in dataset:
16    print(batch_x.numpy(), batch_y.numpy())

Generators are convenient for custom logic, but file-based pipelines and native TensorFlow readers often scale better for production training.

Performance Features That Matter

Three methods provide the biggest practical win:

• 'cache() keeps a dataset in memory or on local storage after the first pass'
• 'prefetch() overlaps input preparation with model execution'
• parallel map() improves preprocessing throughput

For example:

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

This is often enough to remove input bottlenecks in small and medium training jobs.

Common Pitfalls

One common mistake is forgetting to batch the dataset. A model may still run on single examples, but training will be slow and shapes may not match the network’s expectations.

Another issue is shuffling after batching instead of before batching. If you batch first, the model still sees groups that reflect the original order.

People also use the term “dataloader” loosely and then search for a PyTorch-style object in TensorFlow. In TensorFlow 2.x, the normal answer is to build a good tf.data.Dataset pipeline instead.

Finally, be careful with Python generators for high-throughput training. They work, but they can become the slowest part of the system if parsing and augmentation stay in Python rather than moving into TensorFlow operations.

Summary

• In TensorFlow 2.x, tf.data.Dataset is the main data loading pipeline.
• 'from_tensor_slices is the easiest way to create a dataset from in-memory arrays.'
• 'map, shuffle, batch, and prefetch are the core building blocks.'
• Keras can train directly from a dataset with model.fit(dataset).
• For larger workloads, prefer scalable input pipelines over ad hoc Python loading code.