Cosine Annealing LR

Problem

Cosine annealing schedule: lr = min_lr + 0.5 * (base_lr - min_lr) * (1 + cos(pi * step / total_steps)). Clamp step to total_steps before computing.

The signature is cosine_lr(step, total_steps, base_lr, min_lr=0.0).

Hint

At step 0 the cosine term is 1, so the rate starts at base_lr and ends at min_lr.

case 1

in: [0,100,0.001]

expected: 0.001

case 2

in: [50,100,0.001]

expected: 0.0005

case 3

in: [100,100,0.001]

expected: 0

case 4

in: [100,100,0.001,0.0001]

expected: 0.0001

case 5

in: [0,100,0.001,0.0001]

expected: 0.001

Starter codePython

import math
def cosine_lr(step, total_steps, base_lr, min_lr=0.0):
    # Your code here
    pass

5 test cases · run your code against them in the editor