stlr.py

'STLR scheduler from https://arxiv.org/abs/1801.06146'

class STLR(_LRScheduler):
	def __init__(self, optimizer, T_max, last_epoch=-1, ratio=32):
		self.T_max = T_max
		self.cut =  np.floor(T_max*0.1)
		self.ratio = ratio
		super(STLR, self).__init__(optimizer, last_epoch)

	def get_lr(self):
		if self.last_epoch < self.cut:
			p = self.last_epoch/self.cut
		else:
			fraction = (self.last_epoch - self.cut)/(self.cut*(1/self.ratio - 1))
			p = 1 - fraction
		return [base_lr*(1 + p*(self.ratio - 1))/self.ratio for base_lr in self.base_lrs]