Cosine Restarts Learning Rate

Problem

Compute the learning rate for cosine annealing with warm restarts.

Let t = step % period and return min_lr + 0.5 * (base_lr - min_lr) * (1 + cos(pi * t / period)). If period is not positive, return base_lr.

Hint

The modulo resets the cosine cycle every period steps.

case 1

in: [0.1,0.001,10,0]

expected: 0.1

case 2

in: [0.1,0.001,10,5]

expected: 0.0505

case 3

in: [0.1,0.001,10,10]

expected: 0.1

case 4

in: [0.1,0.001,10,15]

expected: 0.0505

case 5

in: [0.1,0.001,0,3]

expected: 0.1

Starter codePython

def cosine_restarts_lr(base_lr, min_lr, period, step):
    # Your code here
    pass

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