Support early stopping during training inside of the early stopping callback

[ananthsub]

🚀 Feature

The ability to early stop during training, as controlled by the callback instead of the training loop.

Motivation

This simplifies the training loop from manually running these callback hooks if the validation epoch won't be run: https://github.com/PyTorchLightning/pytorch-lightning/blob/5bd3cd5f712b65d38812b27cf957261bb06b83c5/pytorch_lightning/trainer/training_loop.py#L152-L159

Similar to flexibility with checkpointing callbacks, this can eventually enable users to specify separate early stop criteria for both training and validation.

Pitch

Add an check_on_train_epoch_end flag to the callback constructor. See #6944 for a sketch

This flag controls whether we check during training or validation. Because the monitor metric may be in training but not validation, or vice versa, this flag makes the check across these two hooks mutually exclusive.

    def on_train_epoch_end(self, trainer, pl_module, outputs) -> None:
        if not self._on_train_epoch_end or self._should_skip_check(trainer):
            return
        self._run_early_stopping_check(trainer)

    def on_validation_end(self, trainer, pl_module):
        if self._on_train_epoch_end or self._should_skip_check(trainer):
            return
        self._run_early_stopping_check(trainer)

For parity with existing behavior, by default this flag will be False by default. With this feature enabled, users can specify their callback for early stopping during training as such:

stop = EarlyStopping(monitor='abc', min_delta=0.1, patience=0, check_on_train_epoch_end=True)

Users could then create multiple such callbacks:

train_stop = EarlyStopping(monitor='abc', min_delta=0.1, patience=0, check_on_train_epoch_end=True)
val_stop = EarlyStopping(monitor='val_loss', min_delta=0.5, patience=3)
trainer = Trainer(...., callbacks=[train_stop, val_stop], ...)
trainer.fit(...)

Alternatives

Keep as is

Additional context