keras
model.fit
initializable iterator
tf.Dataset
TensorFlow

keras model.fit fed with initializable iterator of tf.Dataset object

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

This issue usually comes from mixing TensorFlow 1 dataset patterns with modern Keras training. In current TensorFlow, model.fit normally expects a tf.data.Dataset directly, not a manually initialized iterator, so the usual fix is to pass the dataset itself and let Keras drive iteration.

The TensorFlow 1 Pattern Versus the TensorFlow 2 Pattern

In TensorFlow 1 graph mode, it was common to create an iterator explicitly and initialize it in a session. In TensorFlow 2, Keras handles dataset consumption for you.

The preferred modern pattern is:

python
1import tensorflow as tf
2
3x = tf.random.normal((1000, 20))
4y = tf.cast(tf.reduce_sum(x, axis=1) > 0, tf.float32)
5
6ds = tf.data.Dataset.from_tensor_slices((x, y))
7ds = ds.shuffle(1000).batch(32).prefetch(tf.data.AUTOTUNE)
8
9model = tf.keras.Sequential([
10    tf.keras.layers.Input(shape=(20,)),
11    tf.keras.layers.Dense(32, activation="relu"),
12    tf.keras.layers.Dense(1, activation="sigmoid"),
13])
14
15model.compile(optimizer="adam", loss="binary_crossentropy", metrics=["accuracy"])
16model.fit(ds, epochs=3)

No explicit iterator is needed.

Why Initializable Iterators Feel Wrong with model.fit

An initializable iterator assumes something else will run the initializer before data is consumed. That matches TF1 session-based training, but it does not match the usual Keras fit lifecycle in TensorFlow 2.

Legacy code often looked like this:

python
iterator = tf.compat.v1.data.make_initializable_iterator(dataset)
next_batch = iterator.get_next()

That still belongs to compatibility-mode graph execution, not to the standard TF2 model.fit style.

Repeated Datasets Still Need Step Counts

One place people still get confused is .repeat(). If the dataset repeats forever, Keras no longer knows where an epoch should end, so you must provide steps_per_epoch.

python
1train_ds = ds.repeat()
2
3model.fit(
4    train_ds,
5    epochs=5,
6    steps_per_epoch=200,
7)

That is not an iterator problem. It is just how Keras defines epoch boundaries for potentially infinite datasets.

Check Dataset Element Structure Early

model.fit expects dataset elements to match one of these shapes:

  • '(features, labels)'
  • '(features, labels, sample_weight)'

A quick check before training saves time:

python
for batch in ds.take(1):
    features, labels = batch
    print(features.shape, labels.shape)

If the dataset element structure is wrong, you want to catch that immediately instead of discovering it only after Keras throws a deeper shape error.

What to Do in a Legacy TF1 Codebase

If you are maintaining older TF1 code, keep the iterator logic inside the compatibility boundary rather than spreading it through new training code.

python
1dataset = tf.data.Dataset.from_tensor_slices((x, y)).batch(32)
2iterator = tf.compat.v1.data.make_initializable_iterator(dataset)
3next_features, next_labels = iterator.get_next()
4
5with tf.compat.v1.Session() as sess:
6    sess.run(iterator.initializer)
7    batch_x, batch_y = sess.run([next_features, next_labels])
8    print(batch_x.shape, batch_y.shape)

That is fine as migration glue. It should not be the default design for new Keras pipelines.

Migration Advice for Older Codebases

If your project still contains TF1-era helpers, do the migration in small steps. First replace manual iterator feeding with direct dataset passing. Then remove session-managed training loops. Keeping the compatibility layer narrow is much safer than trying to rewrite the whole input pipeline in one jump.

That incremental approach also makes it easier to compare behavior before and after the change, which is important when older training code already has production history.

Common Pitfalls

  • Trying to feed a TF1-style initializable iterator directly into a normal TF2 model.fit workflow.
  • Using .repeat() without steps_per_epoch, which makes epoch length undefined.
  • Skipping a quick inspection of dataset element structure before training.
  • Mixing session-era code and eager-era Keras patterns in the same training path without a clear boundary.
  • Treating iterator lifecycle as the real problem when the actual issue is the dataset structure or epoch definition.

Summary

  • In modern TensorFlow, pass a tf.data.Dataset directly to model.fit.
  • Initializable iterators are mainly a TensorFlow 1 compatibility concept.
  • Repeated datasets require explicit step counts.
  • Validate dataset element structure before launching training.
  • Keep TF1 iterator code isolated if you still need it during migration.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.