Fix warmup for polynomial decay schedule
Summary: Pull Request resolved:
@@ -62,7 +62,9 @@ class PolynomialDecaySchedule(FairseqLRScheduler):
self.warmup_factor = num_updates / float(self.args.warmup_updates)
self.optimizer.set_lr(self.warmup_factor * self.lr)
else:
- num_updates -= self.args.warmup_updates
- lr = (self.lr - self.end_learning_rate) * (1 - num_updates / self.total_num_update) ** (self.power) + self.end_learning_rate
+ warmup = self.args.warmup_updates
+ lr_range = self.lr - self.end_learning_rate
+ pct_remaining = 1 - (num_updates - warmup) / (self.total_num_update - warmup)
+ lr = lr_range * pct_remaining ** (self.power) + self.end_learning_rate
self.optimizer.set_lr(lr)
return self.optimizer.get_lr()
