Spaces:
Build error
Build error
File size: 1,542 Bytes
222619b |
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 40 41 42 43 44 45 46 47 48 49 50 51 |
from utils.hparams import hparams
class NoneSchedule(object):
def __init__(self, optimizer):
super().__init__()
self.optimizer = optimizer
self.constant_lr = hparams['lr']
self.step(0)
def step(self, num_updates):
self.lr = self.constant_lr
for param_group in self.optimizer.param_groups:
param_group['lr'] = self.lr
return self.lr
def get_lr(self):
return self.optimizer.param_groups[0]['lr']
def get_last_lr(self):
return self.get_lr()
class RSQRTSchedule(object):
def __init__(self, optimizer):
super().__init__()
self.optimizer = optimizer
self.constant_lr = hparams['lr']
self.warmup_updates = hparams['warmup_updates']
self.hidden_size = hparams['hidden_size']
self.lr = hparams['lr']
for param_group in optimizer.param_groups:
param_group['lr'] = self.lr
self.step(0)
def step(self, num_updates):
constant_lr = self.constant_lr
warmup = min(num_updates / self.warmup_updates, 1.0)
rsqrt_decay = max(self.warmup_updates, num_updates) ** -0.5
rsqrt_hidden = self.hidden_size ** -0.5
self.lr = max(constant_lr * warmup * rsqrt_decay * rsqrt_hidden, 1e-7)
for param_group in self.optimizer.param_groups:
param_group['lr'] = self.lr
return self.lr
def get_lr(self):
return self.optimizer.param_groups[0]['lr']
def get_last_lr(self):
return self.get_lr()
|