Progress bar now handles steps_per_execution. (keras-team#21422)

Progress bar would always report the starting batch + 1 at the end of the batch. Now it takes into account `steps_per_execution` for the last batch reported.

Fixes keras-team#20861

Author: hertschuh

keras/src/backend/jax/trainer.py

@@ -408,9 +408,9 @@ def fit(
 
                 self._jax_state_synced = True
                 with epoch_iterator.catch_stop_iteration():
-                    for step, iterator in epoch_iterator:
+                    for begin_step, end_step, iterator in epoch_iterator:
                         # Callbacks
-                        callbacks.on_train_batch_begin(step)
+                        callbacks.on_train_batch_begin(begin_step)
 
                         # Train step
                         if self._jax_state_synced:
@@ -441,7 +441,7 @@ def fit(
                             "metrics_variables": metrics_variables,
                         }
                         # Dispatch callbacks. This takes care of async dispatch.
-                        callbacks.on_train_batch_end(step, logs)
+                        callbacks.on_train_batch_end(end_step, logs)
 
                         if self.stop_training:
                             # Stop training if a callback has set
@@ -569,8 +569,8 @@ def evaluate(
 
         self._jax_state_synced = True
         with epoch_iterator.catch_stop_iteration():
-            for step, iterator in epoch_iterator:
-                callbacks.on_test_batch_begin(step)
+            for begin_step, end_step, iterator in epoch_iterator:
+                callbacks.on_test_batch_begin(begin_step)
 
                 if self._jax_state_synced:
                     # The state may have been synced by a callback.
@@ -600,7 +600,7 @@ def evaluate(
                 }
 
                 # Dispatch callbacks. This takes care of async dispatch.
-                callbacks.on_test_batch_end(step, logs)
+                callbacks.on_test_batch_end(end_step, logs)
 
                 if self.stop_evaluating:
                     break
@@ -633,7 +633,7 @@ def predict(
 
         if not all(layer.built for layer in self._flatten_layers()):
             # Build the model on one batch of data.
-            for _, iterator in epoch_iterator:
+            for _, _, iterator in epoch_iterator:
                 # Build model
                 x, _, _ = data_adapter_utils.unpack_x_y_sample_weight(
                     next(iterator)
@@ -677,8 +677,8 @@ def append_to_outputs(batch_outputs, outputs):
         outputs = None
         non_trainable_variables = None
         with epoch_iterator.catch_stop_iteration():
-            for step, iterator in epoch_iterator:
-                callbacks.on_predict_batch_begin(step)
+            for begin_step, end_step, iterator in epoch_iterator:
+                callbacks.on_predict_batch_begin(begin_step)
                 if self._jax_state_synced:
                     # The state may have been synced by a callback.
                     state = self._get_jax_state(
@@ -701,7 +701,9 @@ def append_to_outputs(batch_outputs, outputs):
                 outputs = append_to_outputs(batch_outputs, outputs)
 
                 # Dispatch callbacks. This takes care of async dispatch.
-                callbacks.on_predict_batch_end(step, {"outputs": batch_outputs})
+                callbacks.on_predict_batch_end(
+                    end_step, {"outputs": batch_outputs}
+                )
 
                 if self.stop_predicting:
                     break

keras/src/backend/numpy/trainer.py

@@ -211,11 +211,11 @@ def append_to_outputs(batch_outputs, outputs):
         self.stop_predicting = False
         callbacks.on_predict_begin()
         outputs = None
-        for step, data in epoch_iterator:
-            callbacks.on_predict_batch_begin(step)
+        for begin_step, end_step, data in epoch_iterator:
+            callbacks.on_predict_batch_begin(begin_step)
             batch_outputs = self.predict_function(data)
             outputs = append_to_outputs(batch_outputs, outputs)
-            callbacks.on_predict_batch_end(step, {"outputs": batch_outputs})
+            callbacks.on_predict_batch_end(end_step, {"outputs": batch_outputs})
             if self.stop_predicting:
                 break
         callbacks.on_predict_end()
@@ -255,7 +255,7 @@ def evaluate(
 
         if not all(layer.built for layer in self._flatten_layers()):
             # Build the model on one batch of data.
-            for _, data in epoch_iterator:
+            for _, _, data in epoch_iterator:
                 data_batch = data[0]
                 self._symbolic_build(data_batch)
                 break
@@ -276,10 +276,10 @@ def evaluate(
         callbacks.on_test_begin()
         logs = {}
         self.reset_metrics()
-        for step, data in epoch_iterator:
-            callbacks.on_test_batch_begin(step)
+        for begin_step, end_step, data in epoch_iterator:
+            callbacks.on_test_batch_begin(begin_step)
             logs = self.test_function(data)
-            callbacks.on_test_batch_end(step, logs)
+            callbacks.on_test_batch_end(end_step, logs)
             if self.stop_evaluating:
                 break
         logs = self._get_metrics_result_or_logs(logs)

keras/src/backend/openvino/trainer.py

@@ -213,11 +213,11 @@ def append_to_outputs(batch_outputs, outputs):
         self.stop_predicting = False
         callbacks.on_predict_begin()
         outputs = None
-        for step, data in epoch_iterator.enumerate_epoch():
-            callbacks.on_predict_batch_begin(step)
+        for begin_step, end_step, data in epoch_iterator.enumerate_epoch():
+            callbacks.on_predict_batch_begin(begin_step)
             batch_outputs = self.predict_function(data)
             outputs = append_to_outputs(batch_outputs, outputs)
-            callbacks.on_predict_batch_end(step, {"outputs": batch_outputs})
+            callbacks.on_predict_batch_end(end_step, {"outputs": batch_outputs})
             if self.stop_predicting:
                 break
         callbacks.on_predict_end()

keras/src/backend/tensorflow/distribute_test.py

@@ -104,7 +104,7 @@ def test_epoch_iterator(self):
             distribute_strategy=strategy,
         )
         steps_seen = []
-        for step, data_iterator in epoch_iterator:
+        for step, _, data_iterator in epoch_iterator:
             steps_seen.append(step)
             batch = next(data_iterator)
             self.assertEqual(len(batch), 3)

keras/src/backend/tensorflow/trainer.py

@@ -372,10 +372,10 @@ def fit(
             self.reset_metrics()
             callbacks.on_epoch_begin(epoch)
             with epoch_iterator.catch_stop_iteration():
-                for step, iterator in epoch_iterator:
-                    callbacks.on_train_batch_begin(step)
+                for begin_step, end_step, iterator in epoch_iterator:
+                    callbacks.on_train_batch_begin(begin_step)
                     logs = self.train_function(iterator)
-                    callbacks.on_train_batch_end(step, logs)
+                    callbacks.on_train_batch_end(end_step, logs)
                     if self.stop_training:
                         break
 
@@ -484,10 +484,10 @@ def evaluate(
         logs = {}
         self.reset_metrics()
         with epoch_iterator.catch_stop_iteration():
-            for step, iterator in epoch_iterator:
-                callbacks.on_test_batch_begin(step)
+            for begin_step, end_step, iterator in epoch_iterator:
+                callbacks.on_test_batch_begin(begin_step)
                 logs = self.test_function(iterator)
-                callbacks.on_test_batch_end(step, logs)
+                callbacks.on_test_batch_end(end_step, logs)
                 if self.stop_evaluating:
                     break
         logs = self._get_metrics_result_or_logs(logs)
@@ -560,12 +560,14 @@ def get_data(iterator):
         callbacks.on_predict_begin()
         outputs = None
         with epoch_iterator.catch_stop_iteration():
-            for step, iterator in epoch_iterator:
-                callbacks.on_predict_batch_begin(step)
+            for begin_step, end_step, iterator in epoch_iterator:
+                callbacks.on_predict_batch_begin(begin_step)
                 data = get_data(iterator)
                 batch_outputs = self.predict_function(data)
                 outputs = append_to_outputs(batch_outputs, outputs)
-                callbacks.on_predict_batch_end(step, {"outputs": batch_outputs})
+                callbacks.on_predict_batch_end(
+                    end_step, {"outputs": batch_outputs}
+                )
                 if self.stop_predicting:
                     break
         callbacks.on_predict_end()
@@ -696,7 +698,7 @@ def _maybe_symbolic_build(self, iterator=None, data_batch=None):
         # Unlike jax/torch iterator, tf iterator returns an iterator instead
         # of data batch in `iterator`.
         if iterator is not None:
-            for _, it in iterator:
+            for _, _, it in iterator:
                 maybe_distributed_data_batch = next(it)
                 has_distributed_values = tree.map_structure(
                     lambda x: isinstance(x, tf.distribute.DistributedValues),

keras/src/backend/torch/trainer.py

@@ -256,14 +256,14 @@ def fit(
             self.train()
 
             logs = {}
-            for step, data in epoch_iterator:
+            for begin_step, end_step, data in epoch_iterator:
                 # Callbacks
-                callbacks.on_train_batch_begin(step)
+                callbacks.on_train_batch_begin(begin_step)
 
                 logs = self.train_function(data)
 
                 # Callbacks
-                callbacks.on_train_batch_end(step, logs)
+                callbacks.on_train_batch_end(end_step, logs)
                 if self.stop_training:
                     break
 
@@ -374,10 +374,10 @@ def evaluate(
         callbacks.on_test_begin()
         logs = {}
         self.reset_metrics()
-        for step, data in epoch_iterator:
-            callbacks.on_test_batch_begin(step)
+        for begin_step, end_step, data in epoch_iterator:
+            callbacks.on_test_batch_begin(begin_step)
             logs = self.test_function(data)
-            callbacks.on_test_batch_end(step, logs)
+            callbacks.on_test_batch_end(end_step, logs)
             if self.stop_evaluating:
                 break
         logs = self._get_metrics_result_or_logs(logs)
@@ -433,11 +433,11 @@ def append_to_outputs(batch_outputs, outputs):
         self.stop_predicting = False
         callbacks.on_predict_begin()
         outputs = None
-        for step, data in epoch_iterator:
-            callbacks.on_predict_batch_begin(step)
+        for begin_step, end_step, data in epoch_iterator:
+            callbacks.on_predict_batch_begin(begin_step)
             batch_outputs = self.predict_function(data)
             outputs = append_to_outputs(batch_outputs, outputs)
-            callbacks.on_predict_batch_end(step, {"outputs": batch_outputs})
+            callbacks.on_predict_batch_end(end_step, {"outputs": batch_outputs})
             if self.stop_predicting:
                 break
         callbacks.on_predict_end()

keras/src/trainers/epoch_iterator.py

@@ -116,7 +116,11 @@ def _enumerate_iterator(self):
                         self._interrupted_warning()
                     break
                 self._steps_seen += self.steps_per_execution
-                yield step, self._current_iterator
+                yield (
+                    step,
+                    step + self.steps_per_execution - 1,
+                    self._current_iterator,
+                )
             if self._num_batches and self._steps_seen >= self._num_batches:
                 self._current_iterator = iter(self._get_iterator())
                 self._steps_seen = 0
@@ -126,7 +130,7 @@ def _enumerate_iterator(self):
             while True:
                 step += self.steps_per_execution
                 self._steps_seen = step + self.steps_per_execution
-                yield step, iterator
+                yield step, step + self.steps_per_execution - 1, iterator
         self.data_adapter.on_epoch_end()
 
     def __iter__(self):
@@ -135,19 +139,19 @@ def __iter__(self):
 
     def __next__(self):
         buffer = []
-        step, iterator = next(self._epoch_iterator)
+        begin_step, end_step, iterator = next(self._epoch_iterator)
         with self.catch_stop_iteration():
             for _ in range(self.steps_per_execution):
                 data = next(iterator)
                 buffer.append(data)
-            return step, buffer
+            return begin_step, end_step, buffer
         if buffer:
-            return step, buffer
+            return begin_step, end_step, buffer
         raise StopIteration
 
     def enumerate_epoch(self):
-        for step, data in self:
-            yield step, data
+        for begin_step, end_step, data in self:
+            yield begin_step, end_step, data
 
     @contextlib.contextmanager
     def catch_stop_iteration(self):

keras/src/trainers/epoch_iterator_test.py

@@ -31,9 +31,10 @@ def test_basic_flow(self, call_type):
             generator = iterator
         else:
             generator = iterator.enumerate_epoch()
-        for step, batch in generator:
+        for begin_step, end_step, batch in generator:
             batch = batch[0]
-            steps_seen.append(step)
+            steps_seen.append(begin_step)
+            self.assertEqual(begin_step, end_step)
             self.assertEqual(len(batch), 3)
             self.assertIsInstance(batch[0], np.ndarray)
         self.assertEqual(steps_seen, [0, 1, 2, 3, 4, 5, 6])
@@ -52,7 +53,7 @@ def test_insufficient_data(self):
         )
         steps_seen = []
         with pytest.warns(match="Your input ran out of data"):
-            for step, _ in iterator:
+            for step, _, _ in iterator:
                 steps_seen.append(step)
         self.assertLen(steps_seen, steps_per_epoch - 2)
 
@@ -96,7 +97,7 @@ def __getitem__(self, idx):
             torch_dataset, batch_size=8, shuffle=True
         )
         iterator = epoch_iterator.EpochIterator(torch_dataloader)
-        for _, batch in iterator:
+        for _, _, batch in iterator:
             batch = batch[0]
             self.assertEqual(batch[0].shape, (8, 2))
             self.assertEqual(batch[1].shape, (8, 1))
@@ -226,7 +227,7 @@ def on_epoch_end(self):
 
         num_epochs = 5
         for epoch in range(num_epochs):
-            for step, batch in epoch_iter:
+            for _, _, _ in epoch_iter:
                 pass
 
         self.assertAllEqual(ds.tracker, [1, 2] * num_epochs)

keras/src/trainers/trainer.py

@@ -1072,7 +1072,7 @@ def to_symbolic_input(v):
                 )
 
             if data_batch is None:
-                for _, data_or_iterator in iterator:
+                for _, _, data_or_iterator in iterator:
                     if isinstance(data_or_iterator, (list, tuple)):
                         data_batch = data_or_iterator[0]
                     else: