callbacks.py

import torch
from pytorch_lightning.callbacks import EarlyStopping


class EarlyStoppingWithLogging(EarlyStopping):
    """EarlyStopping that also logs the patience counter and best score.

    PyTorch Lightning forbids ``self.log()`` inside ``on_validation_end``, so
    the patience state is written directly to the trainer's logger experiment
    (e.g. the TensorBoard ``SummaryWriter``). Writes are rank-0 only to avoid
    duplicate events under DDP.

    Logged scalars (against ``trainer.global_step``):
        - ``patience_wait_count``: consecutive validations without improvement.
        - ``patience``: the configured limit (for overlaying in TensorBoard).
        - ``best_val_loss``: best monitored value seen so far (when finite).
    """

    def _log_state(self, trainer, pl_module) -> None:
        if self._should_skip_check(trainer) or not trainer.is_global_zero:
            return
        exp = getattr(trainer.logger, "experiment", None)
        if exp is None or not hasattr(exp, "add_scalar"):
            return
        step = trainer.global_step
        exp.add_scalar("patience_wait_count", float(self.wait_count), step)
        exp.add_scalar("patience", float(self.patience), step)
        if bool(torch.isfinite(self.best_score)):
            exp.add_scalar("best_val_loss", float(self.best_score), step)

    def on_validation_end(self, trainer, pl_module) -> None:
        super().on_validation_end(trainer, pl_module)
        if not self._check_on_train_epoch_end:
            self._log_state(trainer, pl_module)

    def on_train_epoch_end(self, trainer, pl_module) -> None:
        super().on_train_epoch_end(trainer, pl_module)
        if self._check_on_train_epoch_end:
            self._log_state(trainer, pl_module)