<reponame>292916808/MolCloze<filename>bert/training/optim_schedule.py<gh_stars>0
# !/usr/bin/env python
# -*- coding: utf-8 -*-
# author: chinshin
# datetime: 2020/4/20 15:41
import numpy as np


class Scheduleoptim(object):

    def __init__(self, optimizer, d_model, n_warmup_steps, grad_clip=None):
        self.optimizer = optimizer
        self.n_warmup_steps = n_warmup_steps
        self.n_current_steps = 0
        self.init_lr = np.power(d_model, -0.5)
        self.grad_clip = grad_clip

    def step_and_update_lr(self):
        # modify: do not update learning rate
        self._update_learning_rate()
        if self.grad_clip is not None:
            self.clip_gradient(self.grad_clip)
        self.optimizer.step()

    def clip_gradient(self, grad_clip):
        """
        Clips gradients computed during backpropagation to avoid explosion of gradients.

        :param optimizer: optimizer with the gradients to be clipped
        :param grad_clip: clip value
        """
        for group in self.optimizer.param_groups:
            for param in group["params"]:
                if param.grad is not None:
                    param.grad.data.clamp_(-grad_clip, grad_clip)

    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):
        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


# !/usr/bin/env python
# -*- coding: utf-8 -*-
# author: chinshin
# datetime: 2020/4/20 15:41
import numpy as np


class Scheduleoptim(object):

    def __init__(self, optimizer, d_model, n_warmup_steps, grad_clip=None):
        self.optimizer = optimizer
        self.n_warmup_steps = n_warmup_steps
        self.n_current_steps = 0
        self.init_lr = np.power(d_model, -0.5)
        self.grad_clip = grad_clip

    def step_and_update_lr(self):
        # modify: do not update learning rate
        self._update_learning_rate()
        if self.grad_clip is not None:
            self.clip_gradient(self.grad_clip)
        self.optimizer.step()

    def clip_gradient(self, grad_clip):
        """
        Clips gradients computed during backpropagation to avoid explosion of gradients.

        :param optimizer: optimizer with the gradients to be clipped
        :param grad_clip: clip value
        """
        for group in self.optimizer.param_groups:
            for param in group["params"]:
                if param.grad is not None:
                    param.grad.data.clamp_(-grad_clip, grad_clip)

    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):
        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
# !/usr/bin/env python # -*- coding: utf-8 -*- # author: chinshin # datetime: 2020/4/20 15:41 # modify: do not update learning rate Clips gradients computed during backpropagation to avoid explosion of gradients.
 
         :param optimizer: optimizer with the gradients to be clipped
         :param grad_clip: clip value