Non cumulative model.predict

keras/src/backend/jax/trainer.py

@@ -638,7 +638,7 @@ def evaluate(
 
     @traceback_utils.filter_traceback
     def predict(
-        self, x, batch_size=None, verbose="auto", steps=None, callbacks=None
+        self, x, batch_size=None, verbose="auto", steps=None, callbacks=None, accumulate=True
     ):
         # Create an iterator that yields batches of input data.
         epoch_iterator = JAXEpochIterator(
@@ -718,7 +718,8 @@ def append_to_outputs(batch_outputs, outputs):
                     # during predict(), but it's allowed.
                     "non_trainable_variables": non_trainable_variables,
                 }
-                outputs = append_to_outputs(batch_outputs, outputs)
+                if accumulate:
+                    outputs = append_to_outputs(batch_outputs, outputs)
 
                 # Dispatch callbacks. This takes care of async dispatch.
                 callbacks.on_predict_batch_end(
@@ -731,7 +732,11 @@ def append_to_outputs(batch_outputs, outputs):
         self.jax_state_sync()
         callbacks.on_predict_end()
         self._jax_state = None
-        return tree.map_structure_up_to(batch_outputs, np.concatenate, outputs)
+        if accumulate:
+            if outputs is None:
+                return None
+            return tree.map_structure_up_to(batch_outputs, np.concatenate, outputs)
+        return outputs
 
     def train_on_batch(
         self,

keras/src/backend/numpy/trainer.py

@@ -170,7 +170,7 @@ def fit(
 
     @traceback_utils.filter_traceback
     def predict(
-        self, x, batch_size=None, verbose="auto", steps=None, callbacks=None
+        self, x, batch_size=None, verbose="auto", steps=None, callbacks=None, accumulate=True
     ):
         # Create an iterator that yields batches of input data.
         epoch_iterator = EpochIterator(
@@ -214,12 +214,17 @@ def append_to_outputs(batch_outputs, outputs):
         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)
+            if accumulate:
+                outputs = append_to_outputs(batch_outputs, outputs)
             callbacks.on_predict_batch_end(end_step, {"outputs": batch_outputs})
             if self.stop_predicting:
                 break
         callbacks.on_predict_end()
-        return tree.map_structure_up_to(batch_outputs, np.concatenate, outputs)
+        if accumulate:
+            if outputs is None:
+                return None
+            return tree.map_structure_up_to(batch_outputs, np.concatenate, outputs)
+        return outputs
 
     @traceback_utils.filter_traceback
     def evaluate(

keras/src/backend/openvino/trainer.py

@@ -171,7 +171,7 @@ def fit(
 
     @traceback_utils.filter_traceback
     def predict(
-        self, x, batch_size=None, verbose="auto", steps=None, callbacks=None
+        self, x, batch_size=None, verbose="auto", steps=None, callbacks=None, accumulate=True
     ):
         # Create an iterator that yields batches of input data.
         epoch_iterator = EpochIterator(
@@ -213,15 +213,22 @@ def append_to_outputs(batch_outputs, outputs):
         self.stop_predicting = False
         callbacks.on_predict_begin()
         outputs = None
-        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(end_step, {"outputs": batch_outputs})
+        for _, iterator in epoch_iterator:
+            callbacks.on_predict_batch_begin(epoch_iterator.current_step)
+            batch_outputs = self.predict_function(iterator)
+            if accumulate:
+                outputs = append_to_outputs(batch_outputs, outputs)
+            callbacks.on_predict_batch_end(
+                epoch_iterator.current_step, {"outputs": batch_outputs}
+            )
             if self.stop_predicting:
                 break
         callbacks.on_predict_end()
-        return tree.map_structure_up_to(batch_outputs, np.concatenate, outputs)
+        if accumulate:
+            if outputs is None:
+                return None
+            return tree.map_structure_up_to(batch_outputs, np.concatenate, outputs)
+        return outputs
 
     @traceback_utils.filter_traceback
     def evaluate(

keras/src/backend/tensorflow/trainer.py

@@ -521,7 +521,7 @@ def evaluate(
 
     @traceback_utils.filter_traceback
     def predict(
-        self, x, batch_size=None, verbose="auto", steps=None, callbacks=None
+        self, x, batch_size=None, verbose="auto", steps=None, callbacks=None, accumulate=True
     ):
         # Create an iterator that yields batches of input data.
         epoch_iterator = TFEpochIterator(
@@ -586,17 +586,22 @@ def get_data(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)
+                if accumulate:
+                    outputs = append_to_outputs(batch_outputs, outputs)
                 callbacks.on_predict_batch_end(
                     end_step, {"outputs": batch_outputs}
                 )
                 if self.stop_predicting:
                     break
         callbacks.on_predict_end()
-        outputs = tree.map_structure_up_to(
-            batch_outputs, potentially_ragged_concat, outputs
-        )
-        return tree.map_structure(convert_to_np_if_not_ragged, outputs)
+        if accumulate:
+            if outputs is None:
+                return None
+            outputs = tree.map_structure_up_to(
+                batch_outputs, potentially_ragged_concat, outputs
+            )
+            return tree.map_structure(convert_to_np_if_not_ragged, outputs)
+        return outputs
 
     def train_on_batch(
         self,

keras/src/backend/torch/trainer.py

@@ -395,7 +395,7 @@ def evaluate(
 
     @traceback_utils.filter_traceback
     def predict(
-        self, x, batch_size=None, verbose="auto", steps=None, callbacks=None
+        self, x, batch_size=None, verbose="auto", steps=None, callbacks=None, accumulate=True
     ):
         # Create an iterator that yields batches of input data.
         epoch_iterator = TorchEpochIterator(
@@ -442,13 +442,18 @@ def append_to_outputs(batch_outputs, outputs):
         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)
+            if accumulate:
+                outputs = append_to_outputs(batch_outputs, outputs)
             callbacks.on_predict_batch_end(end_step, {"outputs": batch_outputs})
             if self.stop_predicting:
                 break
         callbacks.on_predict_end()
-        outputs = tree.map_structure(backend.convert_to_numpy, outputs)
-        return tree.map_structure_up_to(batch_outputs, np.concatenate, outputs)
+        if accumulate:
+            if outputs is None:
+                return None
+            outputs = tree.map_structure(backend.convert_to_numpy, outputs)
+            return tree.map_structure_up_to(batch_outputs, np.concatenate, outputs)
+        return outputs
 
     def train_on_batch(
         self,

keras/src/trainers/trainer.py

@@ -807,7 +807,7 @@ def evaluate(
         raise NotImplementedError
 
     def predict(
-        self, x, batch_size=None, verbose="auto", steps=None, callbacks=None
+        self, x, batch_size=None, verbose="auto", steps=None, callbacks=None, accumulate=True
     ):
         """Generates output predictions for the input samples.
 
@@ -858,9 +858,13 @@ def predict(
                 repeating dataset, it will run indefinitely.
             callbacks: List of `keras.callbacks.Callback` instances.
                 List of callbacks to apply during prediction.
+            accumulate: Boolean. Whether to accumulate predictions in memory.
+                If `False`, predictions are not returned and must be handled
+                via callbacks to avoid memory issues with large datasets.
+                Defaults to `True`.
 
         Returns:
-            NumPy array(s) of predictions.
+            NumPy array(s) of predictions if `accumulate=True`, otherwise `None`.
         """
         raise NotImplementedError