Introduction

TensorFlow 2.x prefers eager execution, which means you can usually loop over a tf.data.Dataset directly in Python. But some projects still need session-based execution because they are migrating TensorFlow 1.x code, mixing in graph-only APIs, or running under a compatibility layer. In that case, the proper approach is to build a dataset iterator through tf.compat.v1 and consume it inside a Session until TensorFlow raises OutOfRangeError.

Prefer Eager Iteration When You Can

In native TensorFlow 2 code, iteration is simple:

1import tensorflow as tf
2
3dataset = tf.data.Dataset.from_tensor_slices([10, 20, 30]).batch(2)
4
5for batch in dataset:
6    print(batch.numpy())

That is the idiomatic TensorFlow 2 style. If you are not forced into graph mode, stop there. It is easier to read, easier to debug, and matches how most current TensorFlow APIs are designed.

Session-Based Iteration in TensorFlow 2 Compatibility Mode

If you really need a session, disable eager execution before building the graph and create an iterator explicitly.

1import tensorflow as tf
2
3tf.compat.v1.disable_eager_execution()
4
5dataset = tf.data.Dataset.from_tensor_slices([1, 2, 3, 4, 5]).batch(2)
6iterator = tf.compat.v1.data.make_initializable_iterator(dataset)
7next_batch = iterator.get_next()
8
9with tf.compat.v1.Session() as sess:
10    sess.run(iterator.initializer)
11
12    try:
13        while True:
14            batch = sess.run(next_batch)
15            print(batch)
16    except tf.errors.OutOfRangeError:
17        pass

This pattern is the direct replacement for older TensorFlow 1 style dataset loops. get_next() builds a graph operation, and sess.run(next_batch) fetches one batch each time through the loop.

The OutOfRangeError is normal. It is how TensorFlow signals that the dataset is exhausted.

One-Shot Versus Initializable Iterators

TensorFlow offers more than one iterator style. The two most useful in compatibility mode are one-shot and initializable iterators.

A one-shot iterator is convenient when the dataset does not depend on placeholders or graph-time parameters:

1import tensorflow as tf
2
3tf.compat.v1.disable_eager_execution()
4
5dataset = tf.data.Dataset.range(6).map(lambda x: x * 2)
6iterator = tf.compat.v1.data.make_one_shot_iterator(dataset)
7next_item = iterator.get_next()
8
9with tf.compat.v1.Session() as sess:
10    try:
11        while True:
12            print(sess.run(next_item))
13    except tf.errors.OutOfRangeError:
14        pass

An initializable iterator is better when you need to re-run the dataset or initialize it after feeding placeholders. That is why it is usually the safer default in migrated TensorFlow 1 code.

Repeating, Prefetching, and Resetting

A session loop behaves exactly like the dataset pipeline you build. If you add .repeat(), it will not end until the repeat count is exhausted. If you add .prefetch(), TensorFlow can prepare data in advance but the loop logic stays the same.

1import tensorflow as tf
2
3tf.compat.v1.disable_eager_execution()
4
5dataset = tf.data.Dataset.range(4).repeat(2).batch(2).prefetch(1)
6iterator = tf.compat.v1.data.make_initializable_iterator(dataset)
7next_batch = iterator.get_next()
8
9with tf.compat.v1.Session() as sess:
10    sess.run(iterator.initializer)
11
12    try:
13        while True:
14            print(sess.run(next_batch))
15    except tf.errors.OutOfRangeError:
16        print("Dataset finished")

If you want to iterate again, run the iterator initializer again. That is one of the main advantages of the initializable form.

Common Pitfalls

The first pitfall is mixing eager-style iteration with session-style graph execution in the same dataset pipeline. If eager execution is enabled, for batch in dataset is fine. If you need Session, disable eager execution before constructing the graph and stay within the compatibility API.

Another common mistake is forgetting to initialize the iterator. With an initializable iterator, calling sess.run(next_batch) before sess.run(iterator.initializer) will fail immediately.

Developers also sometimes treat OutOfRangeError as a bug. In dataset iteration it usually means success: TensorFlow reached the end of the pipeline. Catch it and use it to terminate the loop cleanly.

Finally, do not keep building new iterators inside a training loop unless you mean to rebuild the graph repeatedly. Build the dataset and iterator once, initialize when needed, and then consume records with sess.run.

Summary

• In TensorFlow 2, direct Python iteration is preferred whenever possible.
• If you must use a session, use tf.compat.v1 iterators and fetch data with sess.run.
• 'make_initializable_iterator is usually the most flexible option for migrated code.'
• Ending the loop with tf.errors.OutOfRangeError is normal for finite datasets.
• Avoid mixing eager and graph-mode styles in the same input pipeline.