CosineAnnealingLR Learning Rate Scheduler

Problem

Implement the CosineAnnealingLR learning rate scheduler. The learning rate follows a cosine curve from an initial value down to a minimum value eta_min over T_max epochs, then optionally restarts.

Solution

import math

def cosine_annealing_lr(initial_lr: float, eta_min: float, T_max: int, epoch: int) -> float:
    return eta_min + 0.5 * (initial_lr - eta_min) * (1 + math.cos(math.pi * epoch / T_max))

Explanation

1. The formula maps the epoch to a cosine curve between 0 and pi over T_max steps.
2. At epoch 0, cos(0) = 1, so LR = initial_lr.
3. At epoch T_max, cos(pi) = -1, so LR = eta_min.
4. The cosine shape provides a smooth warm-down: the rate of decay is slow at the start, fastest in the middle, and slow again near the end.

Complexity

• Time: O(1)
• Space: O(1)