Noam_Scheduler.py

# https://github.com/tugstugi/pytorch-saltnet/blob/master/utils/lr_scheduler.py

from torch.optim.lr_scheduler import _LRScheduler

class Noam_Scheduler(_LRScheduler):
    def __init__(self, optimizer, warmup_steps):
        super(Noam_Scheduler, self).__init__(optimizer)
        self.warmup_steps = warmup_steps

    def get_lr(self):
        last_epoch = max(1, self.last_epoch)

        scale = self.warmup_steps ** 0.5 * min(last_epoch ** (-0.5), last_epoch * self.warmup_steps ** -1.5)
        return [base_lr * scale for base_lr in self.base_lrs]