tf.data
Datasets
eager execution
TensorFlow
machine learning

How can I use tf.data Datasets in eager execution mode?

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Introduction

In TensorFlow 2.x, eager execution is enabled by default, which means operations run immediately and are easier to debug. The tf.data API still matters in eager mode because it gives you a high performance input pipeline for training and inference. You can iterate through datasets naturally in Python while keeping options like parallel mapping, caching, prefetching, and batching. The key idea is that tf.data controls how data moves and transforms, while eager execution controls when ops are evaluated. This article shows practical dataset patterns that work well in eager mode and explains how to avoid the most common pipeline mistakes.

Core Sections

Build datasets directly from tensors and files

Start with simple constructors and keep transformations incremental. For in-memory arrays, use from_tensor_slices.

python
1import tensorflow as tf
2
3features = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
4labels = tf.constant([0, 1, 1])
5
6ds = tf.data.Dataset.from_tensor_slices((features, labels))
7ds = ds.shuffle(buffer_size=100).batch(2).prefetch(tf.data.AUTOTUNE)
8
9for x_batch, y_batch in ds:
10    print(x_batch.shape, y_batch.numpy())

This loop runs eagerly, so printing and debugging are straightforward.

For file based data, use TextLineDataset, TFRecordDataset, or list_files with interleave when you need parallel reads.

Use map for preprocessing and keep it pure

A good map function should take tensors and return tensors without side effects. Keep random augmentation explicit and use TensorFlow ops when possible.

python
1IMG_SIZE = (224, 224)
2
3def preprocess(path, label):
4    image_bytes = tf.io.read_file(path)
5    image = tf.image.decode_jpeg(image_bytes, channels=3)
6    image = tf.image.resize(image, IMG_SIZE)
7    image = tf.cast(image, tf.float32) / 255.0
8    return image, label
9
10paths = tf.constant(["a.jpg", "b.jpg", "c.jpg"])
11labels = tf.constant([0, 1, 0])
12
13ds = tf.data.Dataset.from_tensor_slices((paths, labels))
14ds = ds.map(preprocess, num_parallel_calls=tf.data.AUTOTUNE)
15ds = ds.batch(32).prefetch(tf.data.AUTOTUNE)

Even in eager mode, this pipeline can be compiled and optimized internally.

Integrate with Keras training

model.fit accepts datasets directly. Make sure your dataset yields (features, labels) tuples with shapes the model expects.

python
1model = tf.keras.Sequential([
2    tf.keras.layers.Input(shape=(224, 224, 3)),
3    tf.keras.layers.Conv2D(16, 3, activation="relu"),
4    tf.keras.layers.GlobalAveragePooling2D(),
5    tf.keras.layers.Dense(2, activation="softmax")
6])
7
8model.compile(optimizer="adam",
9              loss="sparse_categorical_crossentropy",
10              metrics=["accuracy"])
11
12model.fit(ds, epochs=5)

If your dataset is infinite due to repeat(), pass steps_per_epoch so training knows when each epoch ends.

Debug and inspect in eager mode

Because execution is immediate, you can inspect samples at any point:

python
sample_x, sample_y = next(iter(ds.unbatch().take(1)))
print(sample_x.shape, sample_y.numpy())

Use .take(n) generously while building pipelines. It is faster than running full epochs to discover input shape issues.

Common Pitfalls

  • Applying repeat() before shuffle() with a tiny buffer, which can reduce data randomness across epochs.
  • Returning Python objects from map instead of tensors, causing shape/type errors later in training.
  • Forgetting to batch the dataset before model.fit, which can make training very slow and memory inefficient.
  • Using blocking Python logic inside map, preventing parallelism and reducing pipeline throughput.
  • Omitting prefetch, which leaves GPU or TPU resources waiting for input data.

Production Readiness Check

Before closing the task, run a short validation loop on representative inputs and one intentional failure case. Confirm that your code path behaves correctly for normal data, empty data, and malformed data. Capture at least one measurable signal such as runtime, memory use, or error rate, then compare it to your baseline so regressions are visible. Keep this check lightweight so it can run in local development and CI without slowing feedback too much. A simple checklist plus one executable smoke test prevents most regressions after refactors and library upgrades.

text
11. Run happy-path example
22. Run edge-case example
33. Run failure-path example
44. Capture one performance or reliability metric
55. Verify output format and error handling

Summary

tf.data works very well with eager execution and should be your default input pipeline layer in TensorFlow 2.x. Build datasets from tensors or files, keep preprocessing functions tensor-native, and compose shuffle, batch, cache, and prefetch in a clear order. Eager mode makes debugging easier, while tf.data preserves performance at scale. If you validate shapes early and monitor throughput, you can get both developer productivity and strong training performance from the same pipeline.


Course illustration
Course illustration

All Rights Reserved.