Transformerの学習率を調整するSchedulerクラスをPyTorchで書いた
学習率のスケジューラをオリジナルで書きたい場合は、_LRSchedulerを継承し、get_lr関数を自作するのが良い。PyTorch本家がそのような実装となっている。
本家は(v1.13.0リリース後に)実装が変わったばかりのよう
関連
min(step**-0.5, step * self.warmup_epochs**-1.5)の部分が学習率調整の本質
step = self.last_epoch + 1
上記の実装はepoch単位で学習率を調整することが前提になっている
学習の初期段階では0に近い係数(scale)が学習率に掛けられるが、徐々に増えていき、warmup_epochsまでエポック数が達すると係数は1で最大となる。 その後は係数が1より小さくなるため、学習率が減衰されていく。