minerva.schedulers.warmup_cosine_annealing

Classes

WarmupCosineAnnealingLR

A custom learning rate scheduler that combines linear warmup with cosine

Module Contents

class minerva.schedulers.warmup_cosine_annealing.WarmupCosineAnnealingLR(optimizer, warmup_epochs, total_epochs, min_lr=0, last_epoch=-1)

Bases: torch.optim.lr_scheduler._LRScheduler

A custom learning rate scheduler that combines linear warmup with cosine annealing. The learning rate increases linearly over the first ‘warmup_epochs’, and then decreases until ‘total_epochs’ following a cosine curve.

Initializes the scheduler.

Parameters

optimizertorch.optim.optimizer.Optimizer

Wrapped optimizer.

warmup_epochsint

Number of epochs for linear warmup.

total_epochsint

Total number of training epochs.

min_lrfloat

Minimum learning rate expected at the end of the cosine annealing.

last_epoch: int

Index of the last epoch. Used for resuming training.

get_lr()

Compute the next learning rate for each of the optimizer’s param_groups.

Returns:

list[float | Tensor]: A list of learning rates for each of the optimizer’s param_groups with the same types as their current group["lr"]s.

Note

If you’re trying to inspect the most recent learning rate, use get_last_lr() instead.

Note

The returned Tensors are copies, and never alias the optimizer’s group["lr"]s.

min_lr = 0

total_epochs

warmup_epochs

Parameters:

• optimizer (torch.optim.optimizer.Optimizer)

• warmup_epochs (int)

• total_epochs (int)

• min_lr (int)

• last_epoch (int)