Cosine Annealing Learning Rate

Problem

Compute the cosine annealing learning rate:

lr = min_lr + 0.5 * (base_lr - min_lr) * (1 + cos(pi * step / total_steps))

Clamp the progress step / total_steps to [0, 1]. If total_steps is not positive, return min_lr.

Hint

At step = total_steps the cosine term is -1, giving exactly min_lr.

case 1

in: [0.1,0,100,0]

expected: 0.1

case 2

in: [0.1,0,100,50]

expected: 0.05

case 3

in: [0.1,0.01,100,100]

expected: 0.01

case 4

in: [0.1,0,100,150]

expected: 0

case 5

in: [0.2,0.05,50,25]

expected: 0.125

Starter codePython

def cosine_annealing_lr(base_lr, min_lr, total_steps, step):
    # Your code here
    pass

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