lr_scheduler

sorix.optim.lr_scheduler ¶

Learning rate schedulers for sorix optimizers.

Schedulers adjust the learning rate of each parameter group in an optimizer following a policy. They do not touch gradients or parameters — they only rewrite optimizer.param_groups[*]['lr'].

A scheduler is advanced with scheduler.step() after optimizer.step(), once per epoch (not once per mini-batch)::

optimizer = sorix.optim.Adam(model.parameters(), lr=1e-3)
scheduler = sorix.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)

for epoch in range(100):
    for X_batch, y_batch in loader:
        optimizer.zero_grad()
        loss = criterion(model(X_batch), y_batch)
        loss.backward()          # computes gradients
        optimizer.step()         # applies them using the current lr
    scheduler.step()             # picks the lr for the next epoch

ReduceLROnPlateau is the exception: it is metric-driven, so it is stepped with the monitored value, scheduler.step(val_loss).

StepLR ¶

StepLR(optimizer, step_size, gamma=0.1, last_epoch=-1)

Bases: _LRScheduler

Decays the learning rate of each parameter group by gamma every step_size epochs, following the staircase schedule

.. math:: \eta_t = \eta_0 \cdot \gamma^{\lfloor t / s \rfloor}

Parameters:

optimizer (Optimizer) –

Wrapped optimizer.
step_size (int) –

Period of learning rate decay. Must be >= 1.
gamma (float, default: 0.1 ) –

Multiplicative factor of learning rate decay. Default: 0.1.
last_epoch (int, default: -1 ) –

Index of the last completed epoch. Default: -1.

Example::

scheduler = StepLR(optimizer, step_size=30, gamma=0.1)
# lr decays by 0.1× every 30 epochs

Source code in sorix/optim/lr_scheduler.py

def __init__(
    self,
    optimizer: "Optimizer",
    step_size: int,
    gamma: float = 0.1,
    last_epoch: int = -1,
) -> None:
    if step_size < 1:
        raise ValueError(f"step_size must be >= 1, got {step_size}")
    if gamma <= 0.0:
        raise ValueError(f"gamma must be > 0, got {gamma}")
    self.step_size = step_size
    self.gamma = gamma
    super().__init__(optimizer, last_epoch)

ExponentialLR ¶

ExponentialLR(optimizer, gamma, last_epoch=-1)

Bases: _LRScheduler

Decays the learning rate of each parameter group by gamma every epoch, following

.. math:: \eta_t = \eta_0 \cdot \gamma^{t}

Parameters:

optimizer (Optimizer) –

Wrapped optimizer.
gamma (float) –

Multiplicative factor of learning rate decay.
last_epoch (int, default: -1 ) –

Index of the last completed epoch. Default: -1.

Example::

scheduler = ExponentialLR(optimizer, gamma=0.95)
# lr is multiplied by 0.95 each epoch

Source code in sorix/optim/lr_scheduler.py

def __init__(
    self,
    optimizer: "Optimizer",
    gamma: float,
    last_epoch: int = -1,
) -> None:
    if gamma <= 0.0:
        raise ValueError(f"gamma must be > 0, got {gamma}")
    self.gamma = gamma
    super().__init__(optimizer, last_epoch)

CosineAnnealingLR ¶

CosineAnnealingLR(
    optimizer, T_max, eta_min=0.0, last_epoch=-1
)

Bases: _LRScheduler

Anneals the learning rate along a half cosine over T_max epochs:

.. math::

\eta_t = \eta_{\min}
        + \tfrac{1}{2}(\eta_0 - \eta_{\min})
          \left(1 + \cos\left(\frac{\pi t}{T_{\max}}\right)\right)

So lr goes from base_lr at t = 0 down to eta_min at t = T_max.

Note

The formula is periodic with period 2 * T_max. Stepping past T_max makes the learning rate rise back towards base_lr (a "warm restart"). If you train for more than T_max epochs and do not want that, stop stepping the scheduler at T_max.

Parameters:

optimizer (Optimizer) –

Wrapped optimizer.
T_max (int) –

Maximum number of iterations (half-period of the cosine). Must be >= 1.
eta_min (float, default: 0.0 ) –

Minimum learning rate. Default: 0.
last_epoch (int, default: -1 ) –

Index of the last completed epoch. Default: -1.

Example::

scheduler = CosineAnnealingLR(optimizer, T_max=100, eta_min=1e-6)

Source code in sorix/optim/lr_scheduler.py

def __init__(
    self,
    optimizer: "Optimizer",
    T_max: int,
    eta_min: float = 0.0,
    last_epoch: int = -1,
) -> None:
    if T_max < 1:
        raise ValueError(f"T_max must be >= 1, got {T_max}")
    self.T_max = T_max
    self.eta_min = eta_min
    super().__init__(optimizer, last_epoch)

ReduceLROnPlateau ¶

ReduceLROnPlateau(
    optimizer,
    mode="min",
    factor=0.1,
    patience=10,
    min_lr=0.0,
    threshold=0.0001,
)

Reduces learning rate when a metric has stopped improving. Models often benefit from reducing the learning rate by a factor once learning stagnates.

Unlike the epoch-driven schedulers, this one is metric-driven: call step(metric) with the monitored value after each validation pass.

Parameters:

optimizer (Optimizer) –

Wrapped optimizer.
mode (str, default: 'min' ) –

'min' or 'max'. In 'min' mode, lr will be reduced when the quantity monitored has stopped decreasing; in 'max' mode it will be reduced when the quantity has stopped increasing. Default: 'min'.
factor (float, default: 0.1 ) –

Factor by which the learning rate will be reduced. Default: 0.1.
patience (int, default: 10 ) –

Number of epochs with no improvement after which learning rate will be reduced. Default: 10.
min_lr (float, default: 0.0 ) –

A lower bound on the learning rate. Default: 0.
threshold (float, default: 0.0001 ) –

Absolute improvement required to reset the patience counter. Default: 1e-4.

Example::

scheduler = ReduceLROnPlateau(optimizer, mode='min', patience=5, factor=0.5)
for epoch in range(epochs):
    train(...)
    val_loss = validate(...)
    scheduler.step(val_loss)

Source code in sorix/optim/lr_scheduler.py

def __init__(
    self,
    optimizer: "Optimizer",
    mode: str = "min",
    factor: float = 0.1,
    patience: int = 10,
    min_lr: float = 0.0,
    threshold: float = 1e-4,
) -> None:
    if mode not in ("min", "max"):
        raise ValueError(f"mode must be 'min' or 'max', got {mode!r}")
    if factor >= 1.0:
        raise ValueError("factor must be < 1.0")

    self.optimizer = optimizer
    self.mode = mode
    self.factor = factor
    self.patience = patience
    self.min_lr = min_lr
    self.threshold = threshold

    self._best: Optional[float] = None
    self._num_bad_epochs: int = 0

step ¶

step(metrics)

Call after validation with the monitored metric value.

Source code in sorix/optim/lr_scheduler.py

def step(self, metrics: float) -> None:
    """Call after validation with the monitored metric value."""
    if self._is_better(metrics):
        self._best = metrics
        self._num_bad_epochs = 0
    else:
        self._num_bad_epochs += 1

    if self._num_bad_epochs >= self.patience:
        for group in self.optimizer.param_groups:
            new_lr = max(group["lr"] * self.factor, self.min_lr)
            group["lr"] = new_lr
        self._num_bad_epochs = 0

get_last_lr ¶

get_last_lr()

Returns the current learning rate for each parameter group.

Source code in sorix/optim/lr_scheduler.py

def get_last_lr(self) -> List[float]:
    """Returns the current learning rate for each parameter group."""
    return [g["lr"] for g in self.optimizer.param_groups]

state_dict ¶

state_dict()

Returns the state of the scheduler as a dict (excluding the optimizer).

Source code in sorix/optim/lr_scheduler.py

def state_dict(self) -> dict:
    """Returns the state of the scheduler as a dict (excluding the optimizer)."""
    return {k: v for k, v in self.__dict__.items() if k != "optimizer"}

load_state_dict ¶

load_state_dict(state_dict)

Loads the scheduler state.

Note

This scheduler mutates the optimizer's learning rate incrementally, so the restored learning rate is whatever the optimizer currently holds — load the optimizer's own state_dict alongside this one.

Source code in sorix/optim/lr_scheduler.py

def load_state_dict(self, state_dict: dict) -> None:
    """
    Loads the scheduler state.

    Note:
        This scheduler mutates the optimizer's learning rate incrementally,
        so the restored learning rate is whatever the optimizer currently
        holds — load the optimizer's own ``state_dict`` alongside this one.
    """
    state_dict = dict(state_dict)
    state_dict.pop("optimizer", None)
    self.__dict__.update(state_dict)