Spaces:
Sleeping
Sleeping
File size: 928 Bytes
ed00004 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 |
import math
class CosineSchedule:
def __init__(self, min_lr, init_lr, decay_rate, max_epochs) -> None:
self.min_lr = min_lr
self.init_lr = init_lr
self.decay_rate = decay_rate
self.max_epochs = max_epochs
def __call__(self, optimizer, epoch):
"""Decay the learning rate"""
lr = (self.init_lr - self.min_lr) * 0.5 * (
1.0 + math.cos(math.pi * epoch / self.max_epochs)
) + self.min_lr
for param_group in optimizer.param_groups:
param_group["lr"] = lr
class StepSchedule:
def __init__(self, min_lr, init_lr, decay_rate) -> None:
self.min_lr = min_lr
self.init_lr = init_lr
self.decay_rate = decay_rate
def __call__(self, optimizer, epoch):
lr = max(self.min_lr, self.init_lr * (self.decay_rate**epoch))
for param_group in optimizer.param_groups:
param_group["lr"] = lr
|