reduce_lr_on_plateau

reduce_lr_on_plateau

ReduceLROnPlateau callback.

ReduceLROnPlateauCallback

ReduceLROnPlateauCallback(monitor='val_loss', mode='min', factor=0.5, patience=10, threshold=0.0001, threshold_mode='rel', cooldown=0, min_lr=0.0, eps=1e-08)

Bases: Callback

Callback to reduce learning rate when validation loss plateaus.

Wraps PyTorch's ReduceLROnPlateau scheduler to work with Lightning. Reduces learning rate by a factor when validation metric hasn't improved for a given number of epochs (patience).

Initialize ReduceLROnPlateau callback.

Parameters:

Name	Type	Description	Default
monitor	str	Metric to monitor (default: "val_loss")	'val_loss'
mode	str	"min" to reduce LR when metric stops decreasing, "max" for increasing	'min'
factor	float	Factor by which to reduce LR (new_lr = lr * factor)	0.5
patience	int	Number of epochs with no improvement before reducing LR	10
threshold	float	Threshold for measuring improvement	0.0001
threshold_mode	str	"rel" for relative threshold, "abs" for absolute	'rel'
cooldown	int	Number of epochs to wait before resuming normal operation after LR reduction	0
min_lr	float	Minimum learning rate	0.0
eps	float	Minimal decay applied to lr	1e-08

Source code in src/versatil/training/callbacks/reduce_lr_on_plateau.py

def __init__(
    self,
    monitor: str = "val_loss",
    mode: str = "min",
    factor: float = 0.5,
    patience: int = 10,
    threshold: float = 1e-4,
    threshold_mode: str = "rel",
    cooldown: int = 0,
    min_lr: float = 0.0,
    eps: float = 1e-8,
):
    """Initialize ReduceLROnPlateau callback.

    Args:
        monitor: Metric to monitor (default: "val_loss")
        mode: "min" to reduce LR when metric stops decreasing, "max" for increasing
        factor: Factor by which to reduce LR (new_lr = lr * factor)
        patience: Number of epochs with no improvement before reducing LR
        threshold: Threshold for measuring improvement
        threshold_mode: "rel" for relative threshold, "abs" for absolute
        cooldown: Number of epochs to wait before resuming normal operation after LR reduction
        min_lr: Minimum learning rate
        eps: Minimal decay applied to lr
    """
    super().__init__()
    self.monitor = monitor
    self.mode = mode
    self.factor = factor
    self.patience = patience
    self.threshold = threshold
    self.threshold_mode = threshold_mode
    self.cooldown = cooldown
    self.min_lr = min_lr
    self.eps = eps
    self.scheduler: TorchReduceLROnPlateau | None = None
    self._resume_scheduler_state: dict[str, Any] | None = None

on_fit_start

on_fit_start(trainer, pl_module)

Create ReduceLROnPlateau scheduler at start of training.

Restores any scheduler state stashed by load_state_dict so that best, num_bad_epochs, and cooldown_counter survive checkpoint resumes.

Parameters:

Name	Type	Description	Default
trainer	Trainer	Lightning trainer	required
pl_module	LightningModule	Lightning module	required

Source code in src/versatil/training/callbacks/reduce_lr_on_plateau.py

def on_fit_start(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None:
    """Create ReduceLROnPlateau scheduler at start of training.

    Restores any scheduler state stashed by ``load_state_dict`` so that
    ``best``, ``num_bad_epochs``, and ``cooldown_counter`` survive
    checkpoint resumes.

    Args:
        trainer: Lightning trainer
        pl_module: Lightning module
    """
    optimizers = pl_module.optimizers()
    optimizer = optimizers if not isinstance(optimizers, list) else optimizers[0]

    self.scheduler = TorchReduceLROnPlateau(
        optimizer,
        mode=self.mode,
        factor=self.factor,
        patience=self.patience,
        threshold=self.threshold,
        threshold_mode=self.threshold_mode,
        cooldown=self.cooldown,
        min_lr=self.min_lr,
        eps=self.eps,
    )
    if self._resume_scheduler_state is not None:
        self.scheduler.load_state_dict(self._resume_scheduler_state)
        self._resume_scheduler_state = None

state_dict

state_dict()

Return the underlying torch scheduler state for checkpointing.

Returns:

Type	Description
dict[str, Any]	The torch scheduler's state dict, or an empty dict when the
dict[str, Any]	scheduler has not been created yet.

Source code in src/versatil/training/callbacks/reduce_lr_on_plateau.py

def state_dict(self) -> dict[str, Any]:
    """Return the underlying torch scheduler state for checkpointing.

    Returns:
        The torch scheduler's state dict, or an empty dict when the
        scheduler has not been created yet.
    """
    if self.scheduler is None:
        return {}
    return self.scheduler.state_dict()

load_state_dict

load_state_dict(state_dict)

Stash checkpointed scheduler state until the scheduler exists.

on_fit_start recreates the torch scheduler and applies this state onto it, mirroring how the EMA callback handles resume state.

Parameters:

Name	Type	Description	Default
state_dict	dict[str, Any]	Scheduler state saved by state_dict.	required

Source code in src/versatil/training/callbacks/reduce_lr_on_plateau.py

def load_state_dict(self, state_dict: dict[str, Any]) -> None:
    """Stash checkpointed scheduler state until the scheduler exists.

    ``on_fit_start`` recreates the torch scheduler and applies this state
    onto it, mirroring how the EMA callback handles resume state.

    Args:
        state_dict: Scheduler state saved by ``state_dict``.
    """
    if not state_dict:
        return
    self._resume_scheduler_state = state_dict

on_validation_epoch_end

on_validation_epoch_end(trainer, pl_module)

Update scheduler with validation metric at end of epoch.

Parameters:

Name	Type	Description	Default
trainer	Trainer	Lightning trainer	required
pl_module	LightningModule	Lightning module	required

Source code in src/versatil/training/callbacks/reduce_lr_on_plateau.py

def on_validation_epoch_end(
    self, trainer: pl.Trainer, pl_module: pl.LightningModule
) -> None:
    """Update scheduler with validation metric at end of epoch.

    Args:
        trainer: Lightning trainer
        pl_module: Lightning module
    """
    self._step_scheduler(trainer=trainer, pl_module=pl_module)

on_train_epoch_end

on_train_epoch_end(trainer, pl_module)

Update scheduler on train epochs when no validation loop runs.

Runs without a validation dataloader never trigger on_validation_epoch_end, so a train-metric monitor must be stepped from the training loop instead.

Parameters:

Name	Type	Description	Default
trainer	Trainer	Lightning trainer	required
pl_module	LightningModule	Lightning module	required

Source code in src/versatil/training/callbacks/reduce_lr_on_plateau.py

def on_train_epoch_end(
    self, trainer: pl.Trainer, pl_module: pl.LightningModule
) -> None:
    """Update scheduler on train epochs when no validation loop runs.

    Runs without a validation dataloader never trigger
    ``on_validation_epoch_end``, so a train-metric monitor must be
    stepped from the training loop instead.

    Args:
        trainer: Lightning trainer
        pl_module: Lightning module
    """
    if trainer.val_dataloaders is not None:
        return
    self._step_scheduler(trainer=trainer, pl_module=pl_module)