Add terminate_iteration to skip batch handlers

ignite/engine/engine.py

@@ -145,6 +145,7 @@ def __init__(self, process_function: Callable[["Engine", Any], Any]):
         # should_terminate_single_epoch flag: False - don't terminate, True - terminate,
         # "skip_epoch_completed" - terminate and skip the event "EPOCH_COMPLETED"
         self.should_terminate_single_epoch: bool | str = False
+        self.should_terminate_single_iteration: bool = False
         self.should_interrupt = False
         self.state = State()
         self._state_dict_user_keys: list[str] = []
@@ -671,6 +672,19 @@ def terminate_epoch(self, skip_epoch_completed: bool = False) -> None:
         )
         self.should_terminate_single_epoch = "skip_epoch_completed" if skip_epoch_completed else True
 
+    def terminate_iteration(self) -> None:
+        """Signals that the current iteration should finish without firing
+        :attr:`~ignite.engine.events.Events.ITERATION_COMPLETED`.
+
+        This can be used to ignore a batch without triggering handlers that consume ``state.output``.
+        The iteration counter is still incremented, and other events such as
+        :attr:`~ignite.engine.events.Events.ITERATION_STARTED` are still fired.
+
+        .. versionadded:: 0.6.0
+        """
+        self.logger.info("Terminate current iteration is signaled.")
+        self.should_terminate_single_iteration = True
+
     def _handle_exception(self, e: BaseException) -> None:
         if Events.EXCEPTION_RAISED in self._event_handlers:
             self._fire_event(Events.EXCEPTION_RAISED, e)
@@ -986,6 +1000,7 @@ def _internal_run(self) -> State:
 
     def _internal_run_as_gen(self) -> Generator[Any, None, State]:
         self.should_terminate = self.should_terminate_single_epoch = self.should_interrupt = False
+        self.should_terminate_single_iteration = False
         self._init_timers(self.state)
         start_time = time.time()
         try:
@@ -1151,7 +1166,9 @@ def _run_once_on_dataset_as_gen(self) -> Generator[State, None, float]:
                 yield from self._maybe_terminate_or_interrupt()
 
                 self.state.output = self._process_function(self, self.state.batch)
-                self._fire_event(Events.ITERATION_COMPLETED)
+                if not self.should_terminate_single_iteration:
+                    self._fire_event(Events.ITERATION_COMPLETED)
+                self.should_terminate_single_iteration = False
                 yield from self._maybe_terminate_or_interrupt()
 
                 if self.state.epoch_length is not None and iter_counter == self.state.epoch_length:
@@ -1186,6 +1203,7 @@ def _maybe_terminate_legacy(self) -> None:
     def _internal_run_legacy(self) -> State:
         # internal_run without generator for BC
         self.should_terminate = self.should_terminate_single_epoch = self.should_interrupt = False
+        self.should_terminate_single_iteration = False
         self._init_timers(self.state)
         start_time = time.time()
         try:
@@ -1338,7 +1356,9 @@ def _run_once_on_dataset_legacy(self) -> float:
                 self._maybe_terminate_legacy()
 
                 self.state.output = self._process_function(self, self.state.batch)
-                self._fire_event(Events.ITERATION_COMPLETED)
+                if not self.should_terminate_single_iteration:
+                    self._fire_event(Events.ITERATION_COMPLETED)
+                self.should_terminate_single_iteration = False
                 self._maybe_terminate_legacy()
 
                 if self.state.epoch_length is not None and iter_counter == self.state.epoch_length:

tests/ignite/engine/test_engine.py

@@ -9,7 +9,7 @@
 import ignite.distributed as idist
 from ignite.engine import Engine, Events, State
 from ignite.engine.deterministic import keep_random_state
-from ignite.metrics import Average
+from ignite.metrics import Average, RunningAverage
 from tests.ignite.engine import BatchChecker, EpochCounter, IterationCounter
 
 
@@ -52,6 +52,28 @@ def test_terminate(self, skip_completed):
         else:
             assert engine.should_terminate == True  # noqa: E712
 
+    def test_terminate_iteration(self):
+        def process(engine, batch):
+            if batch % 2 == 0:
+                engine.terminate_iteration()
+            return batch
+
+        trainer = Engine(process)
+        completed_iterations = []
+        trainer.add_event_handler(Events.ITERATION_COMPLETED, lambda e: completed_iterations.append(e.state.iteration))
+        RunningAverage(output_transform=lambda output: output).attach(trainer, "running_average")
+
+        state = trainer.run(range(1, 5))
+
+        assert completed_iterations == [1, 3]
+        assert state.metrics["running_average"] == pytest.approx(1.04)
+        assert state.iteration == 4
+        assert not trainer.should_terminate_single_iteration
+
+        evaluator = Engine(process)
+        Average().attach(evaluator, "average")
+        assert evaluator.run(range(1, 5)).metrics["average"] == pytest.approx(2.0)
+
     def test_invalid_process_raises_with_invalid_signature(self):
         with pytest.raises(ValueError, match=r"Engine must be given a processing function in order to run"):
             Engine(None)