Source code for draugr.torch_utilities.optimisation.scheduling.lr_scheduler

from bisect import bisect_right

__all__ = ["WarmupMultiStepLR"]

from torch.optim.lr_scheduler import _LRScheduler


[docs]class WarmupMultiStepLR(_LRScheduler): """description"""
[docs] def __init__( self, optimiser, milestones, gamma=0.1, warmup_factor=1.0 / 3, warmup_iters=500, last_epoch=-1, ): if not list(milestones) == sorted(milestones): raise ValueError( "Milestones should be a list of" " increasing integers. Got {}", milestones, ) self.milestones = milestones self.gamma = gamma self.warmup_factor = warmup_factor self.warmup_iters = warmup_iters super().__init__(optimiser, last_epoch)
[docs] def get_lr(self): """ :return: :rtype:""" warmup_factor = 1 if self.last_epoch < self.warmup_iters: alpha = float(self.last_epoch) / self.warmup_iters warmup_factor = self.warmup_factor * (1 - alpha) + alpha return [ base_lr * warmup_factor * self.gamma ** bisect_right(self.milestones, self.last_epoch) for base_lr in self.base_lrs ]