Adam 优化器在 PyTorch 上进行预热

Question

在论文《Attention is all you need》中，第 5.3 节中，作者建议线性增加学习率，然后按步长平方根倒数按比例减少。

我们如何使用Adam优化器在 PyTorch 中实现这一点？最好没有额外的包。

Answer 1

PyTorch 提供了学习率调度器，用于在训练过程中实现调整学习率的各种方法。一些简单的 LR 调度器已经实现，可以在这里找到： https: //pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate

在您的特殊情况下，您可以 - 就像其他 LR 调度程序一样 - 子类_LRScheduler用于根据时期数实现可变调度。对于一个简单的方法，您只需要实现__init__()和get_lr()方法。

请注意，许多调度程序希望您.step()每个周期调用一次。但您也可以更频繁地更新它，甚至传递自定义参数，就像在余弦退火 LR 调度程序中一样：https://pytorch.org/docs/stable/_modules/torch/optim/lr_scheduler.html#CosineAnnealingLR

Answer 2

正如最后评论中所建议的，我们可以使用https://nlp.seas.harvard.edu/2018/04/03/attention.html#optimizer引入的类。但是这个答案会给出错误，除非我们定义一个函数来更新 state_dict。

这是完整的调度程序：

class NoamOpt:
    "Optim wrapper that implements rate."
    def __init__(self, model_size, warmup, optimizer):
        self.optimizer = optimizer
        self._step = 0
        self.warmup = warmup
        self.model_size = model_size
        self._rate = 0
    
    def state_dict(self):
        """Returns the state of the warmup scheduler as a :class:`dict`.
        It contains an entry for every variable in self.__dict__ which
        is not the optimizer.
        """
        return {key: value for key, value in self.__dict__.items() if key != 'optimizer'}
    
    def load_state_dict(self, state_dict):
        """Loads the warmup scheduler's state.
        Arguments:
            state_dict (dict): warmup scheduler state. Should be an object returned
                from a call to :meth:`state_dict`.
        """
        self.__dict__.update(state_dict) 
        
    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.model_size ** (-0.5) *
            min(step ** (-0.5), step * self.warmup ** (-1.5))) 

稍后，在训练循环中使用它：

optimizer = NoamOpt(input_opts['d_model'], 500,
            torch.optim.Adam(model.parameters(), lr=0, betas=(0.9, 0.98), eps=1e-9))

。。。

optimizer.step()

Answer 3

NoamOpt of Cause 提供了实现预热学习率的路径，如 https://nlp.seas.harvard.edu/2018/04/03/attention.html#optimizer中所示。不过有点旧了，出行不太方便。实现这一点的更聪明的方法是直接使用Pytorch 支持的lambda 学习率调度器。

也就是说，首先定义一个预热函数来自动调整学习率：

def warmup(current_step: int):
if current_step < args.warmup_steps:  # current_step / warmup_steps * base_lr
    return float(current_step / args.warmup_steps)
else:                                 # (num_training_steps - current_step) / (num_training_steps - warmup_steps) * base_lr
    return max(0.0, float(args.training_steps - current_step) / float(max(1, args.training_steps - args.warmup_steps)))

然后构建学习率调度程序并在训练过程中使用它：

lr_scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=warmup)