ScheduledOptim.py

'''A wrapper class for optimizer '''
import numpy as np


class ScheduledOptim():
    '''A simple wrapper class for learning rate scheduling'''

    def __init__(self, optimizer, d_model, max_steps):
        self._optimizer = optimizer
        self.n_warmup_steps = max_steps / 20 
        self.n_current_steps = 0
        self.max_steps = max_steps
        self.init_lr = np.power(d_model, -0.5)

    def step_and_update_lr(self):
        "Step with the inner optimizer"
        self._update_learning_rate()
        self._optimizer.step()

    def zero_grad(self):
        "Zero out the gradients by the inner optimizer"
        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
        if self.n_current_steps <= self.n_warmup_steps:
            lr = self.n_current_steps * 2.4e-4 / self.n_warmup_steps
        else:
            lr = np.power((1 - (self.n_current_steps - self.n_warmup_steps) / (self.max_steps - self.n_warmup_steps)), -0.5) * 2e-4
        for param_group in self._optimizer.param_groups:
            if not param_group['notchange']:
                param_group['lr'] = lr