File size: 1,186 Bytes
9e14838 | 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 31 32 33 34 35 36 37 38 39 | # -*- coding: utf-8 -*-
import torch
from torch.optim import SGD
from torch.optim.lr_scheduler import _LRScheduler
class LinearDecayLR(_LRScheduler):
def __init__(self, optimizer, n_epoch, start_decay, last_epoch=-1, booster=2):
self.start_decay = start_decay
self.n_epoch = n_epoch
self.booster = booster
super(LinearDecayLR, self).__init__(optimizer, last_epoch)
def get_lr(self):
last_epoch = self.last_epoch
n_epoch = self.n_epoch
b_lr = self.base_lrs[-1]
if last_epoch > 0:
try:
cur_lr = self.get_last_lr()
except:
cur_lr = b_lr * self.booster
start_decay = self.start_decay
if last_epoch >= start_decay:
lr = b_lr * self.booster - (b_lr * self.booster) / (
n_epoch - start_decay
) * (last_epoch - start_decay)
else:
if last_epoch < start_decay:
lr = b_lr + (b_lr * self.booster - b_lr) / start_decay * last_epoch
else:
lr = cur_lr
self._last_lr = lr
print(f"Active Learning Rate --- {lr}")
return [lr]
|