TensorFlow
tf.data
Dataset Iteration
Machine Learning
Python

What is the correct way to iterate over an indefinitely repeated tf.data Dataset in Tensorflow 2.0

Master System Design with Codemia

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

Introduction

A tf.data.Dataset created with .repeat() and no count has no natural end. That is useful for training loops, but it also means you must provide your own stopping rule. The correct iteration pattern depends on context: use .take(...) for finite inspection, a bounded Python loop for custom training code, or steps_per_epoch when training with Keras.

What .repeat() Actually Changes

Calling .repeat() without an argument makes the dataset cycle forever.

python
1import tensorflow as tf
2
3dataset = tf.data.Dataset.range(5).repeat()
4
5for value in dataset.take(12):
6    print(int(value))

The output repeats 0 through 4 again and again. Without .take(12), that loop would not terminate.

This is the first important point: an infinite dataset is not a bug. It is a deliberate signal that some other part of the code must decide when to stop.

Iterating In Manual Training Loops

If you are writing your own loop in TensorFlow 2, the clearest pattern is to create an iterator and bound the number of steps explicitly.

python
1import tensorflow as tf
2
3dataset = tf.data.Dataset.range(10).batch(3).repeat()
4iterator = iter(dataset)
5
6for step in range(4):
7    batch = next(iterator)
8    print(f"step={step}, batch={batch.numpy()}")

This keeps the stopping rule visible in plain Python. That is often easier to reason about than a raw for batch in dataset: loop when the dataset is infinite.

You can also bound the pipeline itself with .take(...).

python
1import tensorflow as tf
2
3dataset = tf.data.Dataset.range(10).batch(3).repeat()
4
5for batch in dataset.take(4):
6    print(batch.numpy())

Both forms are valid. .take(...) is concise for quick iteration, while an explicit iterator is often clearer inside custom training code where step counters already exist.

Using model.fit With Infinite Datasets

When training with Keras, the normal pattern is to leave the dataset infinite and tell model.fit how many batches make up one epoch.

python
1import tensorflow as tf
2
3x = tf.random.normal((20, 4))
4y = tf.random.uniform((20,), maxval=2, dtype=tf.int32)
5
6dataset = tf.data.Dataset.from_tensor_slices((x, y))
7dataset = dataset.shuffle(20).batch(4).repeat()
8
9model = tf.keras.Sequential([
10    tf.keras.layers.Input(shape=(4,)),
11    tf.keras.layers.Dense(8, activation="relu"),
12    tf.keras.layers.Dense(2)
13])
14
15model.compile(
16    optimizer="adam",
17    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
18    metrics=["accuracy"]
19)
20
21model.fit(dataset, epochs=3, steps_per_epoch=5)

steps_per_epoch=5 is the stopping rule for each epoch. Without it, Keras cannot infer when an epoch should end because the input never ends.

The same rule applies to validation datasets repeated indefinitely. You must set validation_steps if validation input is infinite.

Choosing The Right Pattern

Use these rules:

  • use .take(n) when you want a finite preview or small evaluation pass
  • use iter(dataset) plus next(...) when you control training steps manually
  • use steps_per_epoch and validation_steps with model.fit

All three are legitimate. The wrong approach is pretending an infinite dataset should behave like a finite one.

Avoiding Hidden Infinite Loops

A common anti-pattern looks like this:

python
for batch in dataset:
    train_step(batch)

That is fine for a finite dataset. It is dangerous for dataset.repeat() because the loop has no stopping condition. If the outer training logic expects epochs, metrics resets, or checkpoint intervals, that assumption is now broken.

The safer version is explicit.

python
for step, batch in enumerate(dataset.take(100)):
    train_step(batch)

Now the loop length is obvious to anyone reading the code.

Common Pitfalls

The most common mistake is calling .repeat() without also defining where iteration should stop. Infinite input requires a finite controller.

Another common issue is forgetting steps_per_epoch when passing an infinite dataset to model.fit. Keras then has no way to determine epoch boundaries.

Developers also mix up dataset repetition and batching semantics. Repeating the dataset does not change batch size or shape. It only restarts the input sequence after exhaustion.

Finally, do not rely on manual interruption such as stopping the notebook cell or killing the process. If the loop should end after a known number of steps, encode that rule directly in the pipeline or loop.

Summary

  • '.repeat() without a count creates an infinite dataset by design.'
  • Use .take(n) for bounded inspection or evaluation.
  • Use an explicit iterator and a step loop for custom training code.
  • Use steps_per_epoch with model.fit when the dataset repeats forever.
  • Add validation_steps if the validation dataset is also infinite.
  • Make stopping conditions explicit instead of relying on external interruption.

Course illustration
Course illustration

All Rights Reserved.