在 Pytorch 中实现自定义学习率调度器？

Question

我想实现这种学习率方法，正如论文中所描述的那样，Attention is all you need。我在 Tensorflow 中有这段代码，但我也想在 Pytorch 中实现它。我知道 Pytorch 有这方面的模块 (https://pytorch.org/docs/stable/_modules/torch/optim/lr_scheduler.html)，但我怎么能着手制作自定义调度程序呢？或者也许上述 lr_scheduler 之一已经实现了相同的功能？

张量流代码：

class CustomSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):
  def __init__(self, d_model, warmup_steps=4000):
    super(CustomSchedule, self).__init__()

    self.d_model = d_model
    self.d_model = tf.cast(self.d_model, tf.float32)

    self.warmup_steps = warmup_steps

  def __call__(self, step):
    arg1 = tf.math.rsqrt(step)
    arg2 = step * (self.warmup_steps ** -1.5)

    return tf.math.rsqrt(self.d_model) * tf.math.minimum(arg1, arg2)

learning_rate = CustomSchedule(d_model)

optimizer = tf.keras.optimizers.Adam(learning_rate, beta_1=0.9, beta_2=0.98,
                                     epsilon=1e-9)

火炬？

import torch 

optimizer = torch.optim.Adam(model.parameters(), lr=0.0001, betas=(0.9, 0.98), eps=1e-9)

scheduler =

Answer 1

由于这是一篇流行论文 (Attention is all you need) 中使用的调度程序，因此网上已经存在相当不错的实现。

您可以从 this repository by @jadore801120 获取 PyTorch 实现。

一旦你拥有它，那么简单

optimizer = torch.optim.Adam(model.parameters(), lr=0.0001, betas=(0.9, 0.98), eps=1e-9)
sched = ScheduledOptim(optimizer, d_model=..., n_warmup_steps=...)

还要确保在正确的时间调用调度程序

for i, batch in enumerate(dataloader):
    sched.zero_grad()
    ...
    loss.backward()
    sched.step_and_update_lr()

在 Pytorch 中实现自定义学习率调度器？

Implementing custom learning rate scheduler in Pytorch?

transformer

tensorflow

attention-model

pytorch