Cosine Decay Learning Rate Scheduler

In deep learning optimization, the learning rate is one of the most critical hyperparameters that determines how quickly or slowly a model converges during training. A static learning rate often leads to suboptimal results—too high causes oscillation and divergence, while too low results in painfully slow convergence.

Learning rate scheduling provides an elegant solution by dynamically adjusting the learning rate throughout training. Among the various scheduling strategies, cosine annealing (also known as cosine decay) has emerged as one of the most effective and widely adopted approaches in modern deep learning.

The Cosine Annealing Formula

The cosine annealing schedule smoothly decreases the learning rate following a cosine curve from the initial learning rate $$\eta_{\text{max}}$$ to a minimum learning rate $$\eta_{\text{min}}$$ over $$T_{\text{max}}$$ epochs:

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

Where:

• $$\eta_t$$ is the learning rate at epoch $$t$$
• $$\eta_{\text{max}}$$ is the initial (maximum) learning rate
• $$\eta_{\text{min}}$$ is the minimum learning rate
• $$T_{\text{max}}$$ is the total number of epochs in one annealing cycle
• $$t$$ is the current epoch number (0-indexed)

Key Properties

1. Smooth Decay: The cosine function provides a gradual transition, starting with slow decay at the beginning, accelerating in the middle, and slowing down again near the end.

2. Bounded Range: The learning rate always remains within $$[\eta_{\text{min}}, \eta_{\text{max}}]$$

3. Natural Convergence: The slow decay near the minimum allows the model to fine-tune and settle into a good local minimum.

Your Task

Implement a Python class CosineDecayScheduler that schedules the learning rate using cosine annealing:

• The __init__ method should accept initial_lr (float), T_max (int), and min_lr (float) as parameters
• The get_lr(self, epoch: int) method should return the learning rate for the given epoch
• The returned learning rate should be rounded to 4 decimal places
• Use only the standard Python library and the math module for trigonometric functions