|
import numpy as np |
|
|
|
|
|
class ScheduledOptim(): |
|
''' A simple wrapper class for learning rate scheduling ''' |
|
|
|
def __init__(self, optimizer, d_model, n_warmup_steps, current_steps): |
|
self._optimizer = optimizer |
|
self.n_warmup_steps = n_warmup_steps |
|
self.n_current_steps = current_steps |
|
self.init_lr = np.power(d_model, -0.5) |
|
|
|
def step_and_update_lr(self): |
|
self._update_learning_rate() |
|
self._optimizer.step() |
|
|
|
def zero_grad(self): |
|
|
|
self._optimizer.zero_grad() |
|
|
|
def _get_lr_scale(self): |
|
return np.min([ |
|
np.power(self.n_current_steps, -0.5), |
|
np.power(self.n_warmup_steps, -1.5) * self.n_current_steps]) |
|
|
|
def _update_learning_rate(self): |
|
''' Learning rate scheduling per step ''' |
|
self.n_current_steps += 1 |
|
lr = self.init_lr * self._get_lr_scale() |
|
|
|
for param_group in self._optimizer.param_groups: |
|
param_group['lr'] = lr |
|
|
