torchcell.scheduler.cosine_annealing_warmup.CosineAnnealingWarmupRestarts

class CosineAnnealingWarmupRestarts(optimizer: Optimizer, first_cycle_steps: int, cycle_mult: float = 1.0, max_lr: float = 0.1, min_lr: float = 0.001, warmup_steps: int = 0, gamma: float = 1.0, last_epoch: int = -1)

Bases: _LRScheduler

Cosine annealing scheduler with linear warmup and decaying cyclic restarts.

optimizer (Optimizer): Wrapped optimizer. first_cycle_steps (int): First cycle step size. cycle_mult(float): Cycle steps magnification. Default: -1. max_lr(float): First cycle’s max learning rate. Default: 0.1. min_lr(float): Min learning rate. Default: 0.001. warmup_steps(int): Linear warmup step size. Default: 0. gamma(float): Decrease rate of max learning rate by cycle. Default: 1. last_epoch (int): The index of last epoch. Default: -1.

init_lr() → None

Set every param group’s learning rate to min_lr and record base LRs.

get_lr() → list[float | Tensor]

Compute the learning rates for the current step (warmup or cosine).

step(epoch: int | None = None) → None

Advance the schedule and update each param group’s learning rate.