Блог

Я хотел бы реализовать этот метод скорости обучения, так как в статье Внимание-это все, что вам нужно. У меня есть этот код в 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 =
 

Поскольку это планировщик, используемый в популярной статье (внимание-это все, что вам нужно), достаточно хорошие реализации уже существуют в Интернете.

Вы можете получить реализацию PyTorch из этого репозитория по @jadore801120.

Как только вы его получите, тогда просто

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