A cosine learning rate scheduler can be implemented from scratch as a custom Python class. The class calculates the decayed learning rate based on the current step according to a cosine curve. It optionally includes a warmup phase where the learning rate increases linearly. Once past the warmup steps but within the maximum update steps, the learning rate is determined by the formula $$\eta_t = \eta_T + \frac{\eta_0 - \eta_T}{2} \left(1 + \cos\left(\frac{\pi (t - t_{\text{warmup}})}{T_{\text{max\_steps}}}\right)\right)$$. The following code demonstrates this implementation and plots the resulting schedule:

```python
class CosineScheduler:
    def __init__(self, max_update, base_lr=0.01, final_lr=0,
               warmup_steps=0, warmup_begin_lr=0):
        self.base_lr_orig = base_lr
        self.max_update = max_update
        self.final_lr = final_lr
        self.warmup_steps = warmup_steps
        self.warmup_begin_lr = warmup_begin_lr
        self.max_steps = self.max_update - self.warmup_steps

    def get_warmup_lr(self, epoch):
        increase = (self.base_lr_orig - self.warmup_begin_lr) \
                       * float(epoch) / float(self.warmup_steps)
        return self.warmup_begin_lr + increase

    def __call__(self, epoch):
        if epoch < self.warmup_steps:
            return self.get_warmup_lr(epoch)
        if epoch <= self.max_update:
            self.base_lr = self.final_lr + (
                self.base_lr_orig - self.final_lr) * (1 + math.cos(
                math.pi * (epoch - self.warmup_steps) / self.max_steps)) / 2
        return self.base_lr

scheduler = CosineScheduler(max_update=20, base_lr=0.3, final_lr=0.01)
d2l.plot(torch.arange(num_epochs), [scheduler(t) for t in range(num_epochs)])
```

Cosine Learning Rate Scheduler Implementation

To empirically demonstrate the impact of a cosine learning rate schedule, a neural network can be trained using an optimizer configured with this scheduler. The schedule dictates that the learning rate decreases gently at the beginning, drops more steeply, and then refines the solution at the end with a very small step size. While this approach can lead to improved training results and final accuracy—particularly in computer vision tasks—such improvements are heuristic and not strictly guaranteed across all models and datasets.

```python
net = net_fn()
trainer = torch.optim.SGD(net.parameters(), lr=0.3)
train(net, train_iter, test_iter, num_epochs, loss, trainer, device,
      scheduler)
```

Cosine Learning Rate Schedule Example

A cosine learning rate schedule, proposed by Loshchilov and Hutter (2016), dynamically adjusts the learning rate by following the shape of a cosine curve. It relies on the observation that the learning rate should not decrease too drastically at the beginning of training, and that the solution should be refined at the end using a very small learning rate. For learning rates in the range $$t \in [0, T]$$, this results in a schedule with the functional form:

$$\eta_t = \eta_T + \frac{\eta_0 - \eta_T}{2} \left(1 + \cos(\pi t/T)\right)$$

Here, $$\eta_0$$ is the initial learning rate and $$\eta_T$$ is the target rate at the maximum update step $$T$$. For steps $$t > T$$, the learning rate is simply pinned to $$\eta_T$$ without increasing it again.

Claude

A learning rate scheduler is a mechanism used to dynamically adjust the learning rate during the training of a model. Rather than keeping the learning rate constant, the scheduler acts as a function that takes the number of optimization updates (such as iterations or epochs) and outputs the appropriate learning rate value for the next step.

Learning Rate Scheduler

Dive into Deep Learning

Applying a learning rate scheduler to gently decrease the learning rate over the course of training can lead to improved model accuracy and less overfitting compared to using a constant learning rate. While the exact cause is debated, one theoretical explanation suggests that taking smaller step sizes forces the model parameters to remain closer to zero, resulting in a simpler model, although this does not completely explain the phenomenon.

Effect of Learning Rate Scheduling on Overfitting

A polynomial learning rate decay schedule gradually reduces the learning rate over the course of training according to a polynomial function. It is considered one of the common policy choices for dynamically adjusting the learning rate.

Polynomial Learning Rate Decay

A piecewise constant learning rate schedule, often referred to as a multi-factor scheduler, maintains a steady learning rate for predefined intervals of training and abruptly drops the rate at specific milestones. Mathematically, given a set of milestone times $$s$$ (such as $$s = \{5, 10, 20\}$$), the learning rate is updated according to the rule $$\eta_{t+1} \leftarrow \eta_t \cdot \alpha$$ whenever the current step $$t \in s$$, where $$\alpha$$ is the designated decay factor.

Piecewise Constant Learning Rate Schedule

Cosine Learning Rate Schedule

To address the dilemma of choosing an initial learning rate that is either too small (causing slow progress) or too large (causing divergence), a simple strategy called optimizer warmup is used. During a warmup period, the learning rate gradually increases—typically linearly—from a small value to its initial maximum, after which it cools down until the end of the optimization process.

Optimizer Warmup

A factor learning rate scheduler is an alternative to polynomial decay that utilizes a multiplicative reduction strategy. The learning rate is updated at each step using the equation $$\eta_{t+1} \leftarrow \eta_t \cdot \alpha$$, where $$\alpha \in (0, 1)$$. To prevent the learning rate from decaying beyond a reasonable lower bound, the update equation is typically modified to $$\eta_{t+1} \leftarrow \mathop{\mathrm{max}}(\eta_{\mathrm{min}}, \eta_t \cdot \alpha)$$.

Factor Learning Rate Scheduler

A basic learning rate adjustment can be performed explicitly at each step by modifying the optimizer's parameter groups directly. The following Python code demonstrates how to set a new learning rate manually in a PyTorch optimizer:

```python
lr = 0.1
trainer.param_groups[0]["lr"] = lr
print(f'learning rate is now {trainer.param_groups[0]["lr"]:.2f}')
```

Explicit Learning Rate Adjustment Implementation

To empirically study learning rate scheduling, a computationally efficient yet nontrivial toy problem can be constructed. A common setup involves training a modernized LeNet architecture on the Fashion-MNIST dataset. This modernized LeNet updates the classic design by replacing sigmoid activations with ReLU activations and substituting AveragePooling operations with MaxPooling.

Learning Rate Scheduler Toy Problem

A simple learning rate scheduler can be defined to decay the learning rate proportionally to the inverse square root of the number of updates. Specifically, at step $$t$$, the learning rate is set to $$\eta = \eta_0 (t + 1)^{-\frac{1}{2}}$$, where $$\eta_0$$ is the initial learning rate. This formulation ensures that the step size gently and continuously decreases as training progresses.

Learn Before

Related

Learn After