File size: 1,392 Bytes
ad9fbbf | 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 | import torch
class NoamOpt:
"Optim wrapper that implements rate."
def __init__(self, model_size, factor, warmup, optimizer):
self.optimizer = optimizer
self._step = 0
self.warmup = warmup
self.factor = factor
self.model_size = model_size
self._rate = 0
def step(self):
"Update parameters and rate"
self._step += 1
rate = self.rate()
for p in self.optimizer.param_groups:
p['lr'] = rate
self._rate = rate
self.optimizer.step()
def rate(self, step = None):
"Implement `lrate` above"
if step is None:
step = self._step
return self.factor * (self.model_size ** (-0.5) * min(step ** (-0.5), step * self.warmup ** (-1.5)))
def zero_grad(self):
self.optimizer.zero_grad()
def get_std_opt(task, parameters, d_model):
train_size = {"PRO":335, "CA":1550, "MG":1729, "MN":547, "Metal":5469}
batch_size = 32
warmup_epoch = (5 if task != "Metal" else 2)
step_each_epoch = int(train_size[task] / batch_size)
warmup = warmup_epoch * step_each_epoch
top_lr = 0.0004
factor = top_lr / (d_model ** (-0.5) * min(warmup ** (-0.5), warmup * warmup ** (-1.5)))
return NoamOpt(
d_model, factor, warmup, torch.optim.Adam(parameters, lr=0, betas=(0.9, 0.98), eps=1e-9)
)
|