Add support for early stopping during training epoch end

Author: ananthsub

pytorch_lightning/callbacks/early_stopping.py

@@ -80,6 +80,7 @@ def __init__(
         verbose: bool = False,
         mode: str = 'min',
         strict: bool = True,
+        during_training: bool = False,
     ):
         super().__init__()
         self.monitor = monitor
@@ -90,6 +91,7 @@ def __init__(
         self.wait_count = 0
         self.stopped_epoch = 0
         self.mode = mode
+        self.during_training = during_training
 
         if self.mode not in self.mode_dict:
             raise MisconfigurationException(f"`mode` can be {', '.join(self.mode_dict.keys())}, got {self.mode}")
@@ -129,15 +131,24 @@ def on_save_checkpoint(self, trainer, pl_module, checkpoint: Dict[str, Any]) ->
             'patience': self.patience
         }
 
-    def on_load_checkpoint(self, callback_state: Dict[str, Any]):
+    def on_load_checkpoint(self, callback_state: Dict[str, Any]) -> None:
         self.wait_count = callback_state['wait_count']
         self.stopped_epoch = callback_state['stopped_epoch']
         self.best_score = callback_state['best_score']
         self.patience = callback_state['patience']
 
-    def on_validation_end(self, trainer, pl_module):
+    def _should_skip_check(self, trainer) -> bool:
         from pytorch_lightning.trainer.states import TrainerState
-        if trainer.state != TrainerState.FITTING or trainer.sanity_checking:
+        return trainer.state != TrainerState.FITTING or trainer.sanity_checking
+
+    def on_train_epoch_end(self, trainer, pl_module, outputs) -> None:
+        if not self.during_training or self._should_skip_check(trainer):
+            return
+        self._run_early_stopping_check(trainer)
+
+
+    def on_validation_end(self, trainer, pl_module):
+        if self.during_training or self._should_skip_check(trainer):
             return
 
         self._run_early_stopping_check(trainer)
@@ -153,7 +164,7 @@ def _run_early_stopping_check(self, trainer):
             trainer.fast_dev_run  # disable early_stopping with fast_dev_run
             or not self._validate_condition_metric(logs)  # short circuit if metric not present
         ):
-            return  # short circuit if metric not present
+            return
 
         current = logs.get(self.monitor)
 

tests/callbacks/test_early_stopping.py

@@ -333,46 +333,77 @@ def test_early_stopping_mode_options():
 
 class EarlyStoppingModel(BoringModel):
 
-    def __init__(self, expected_end_epoch):
+    def __init__(self, expected_end_epoch: int, during_training: bool):
         super().__init__()
         self.expected_end_epoch = expected_end_epoch
+        self.during_training = during_training
+
+    def training_epoch_end(self, outputs):
+        if not self.during_training:
+            return
+        losses = [8, 4, 2, 3, 4, 5, 8, 10]
+        loss = losses[self.current_epoch]
+        self.log('abc', torch.tensor(loss))
+        self.log('cba', torch.tensor(0))
 
     def validation_epoch_end(self, outputs):
+        if self.during_training:
+            return
         losses = [8, 4, 2, 3, 4, 5, 8, 10]
-        val_loss = losses[self.current_epoch]
-        self.log('abc', torch.tensor(val_loss))
+        loss = losses[self.current_epoch]
+        self.log('abc', torch.tensor(loss))
         self.log('cba', torch.tensor(0))
 
     def on_train_end(self) -> None:
         assert self.trainer.current_epoch == self.expected_end_epoch, 'Early Stopping Failed'
 
 
 @pytest.mark.parametrize(
-    "callbacks, expected_stop_epoch, accelerator, num_processes",
+    "callbacks, expected_stop_epoch, during_training, accelerator, num_processes",
     [
-        ([EarlyStopping(monitor='abc'), EarlyStopping(monitor='cba', patience=3)], 3, None, 1),
+        ([EarlyStopping(monitor='abc'), EarlyStopping(monitor='cba', patience=3)], 3, False, None, 1),
         ([EarlyStopping(monitor='cba', patience=3),
-          EarlyStopping(monitor='abc')], 3, None, 1),
+          EarlyStopping(monitor='abc')], 3, False, None, 1),
         pytest.param([EarlyStopping(monitor='abc'),
                       EarlyStopping(monitor='cba', patience=3)],
                      3,
+                     False,
                      'ddp_cpu',
                      2,
                      marks=RunIf(skip_windows=True)),
         pytest.param([EarlyStopping(monitor='cba', patience=3),
                       EarlyStopping(monitor='abc')],
                      3,
+                     False,
+                     'ddp_cpu',
+                     2,
+                     marks=RunIf(skip_windows=True)),
+        ([EarlyStopping(monitor='abc', during_training=True), EarlyStopping(monitor='cba', patience=3, during_training=True)], 3, True, None, 1),
+        ([EarlyStopping(monitor='cba', patience=3, during_training=True),
+          EarlyStopping(monitor='abc', during_training=True)], 3, True, None, 1),
+        pytest.param([EarlyStopping(monitor='abc', during_training=True),
+                      EarlyStopping(monitor='cba', patience=3, during_training=True)],
+                     3,
+                     True,
+                     'ddp_cpu',
+                     2,
+                     marks=RunIf(skip_windows=True)),
+        pytest.param([EarlyStopping(monitor='cba', patience=3, during_training=True),
+                      EarlyStopping(monitor='abc', during_training=True)],
+                     3,
+                     True,
                      'ddp_cpu',
                      2,
                      marks=RunIf(skip_windows=True)),
+
     ],
 )
 def test_multiple_early_stopping_callbacks(
-    tmpdir, callbacks: List[EarlyStopping], expected_stop_epoch: int, accelerator: Optional[str], num_processes: int
+    tmpdir, callbacks: List[EarlyStopping], expected_stop_epoch: int, during_training: bool, accelerator: Optional[str], num_processes: int
 ):
     """Ensure when using multiple early stopping callbacks we stop if any signals we should stop."""
 
-    model = EarlyStoppingModel(expected_stop_epoch)
+    model = EarlyStoppingModel(expected_stop_epoch, during_training)
 
     trainer = Trainer(
         default_root_dir=tmpdir,