Transformerの学習率を調整するSchedulerクラスをPyTorchで書いた
https://tam5917.hatenablog.com/entry/2022/09/16/214700
Attention is All You Needという論文で「warmup & ステップ数の逆平方根で学習率を減衰」させる学習率スケジューリングが提案されたが、そのようなスケジューリングを手軽に行うスケジューラを書いたということである。
学習率のスケジューラをオリジナルで書きたい場合は、_LRSchedulerを継承し、get_lr関数を自作するのが良い。PyTorch本家がそのような実装となっている。
ブログ執筆時の実装 https://github.com/pytorch/pytorch/blob/v1.12.1/torch/optim/lr_scheduler.py#L171
本家は(v1.13.0リリース後に)実装が変わったばかりのよう
https://github.com/pytorch/pytorch/blob/5f09a6d573a2a07c00c76c3cbdbffe0fafe2436d/torch/optim/lr_scheduler.py#L185
関連
https://github.com/pytorch/pytorch/pull/88503
https://github.com/pytorch/pytorch/issues/61232
min(step**-0.5, step * self.warmup_epochs**-1.5)の部分が学習率調整の本質
step = self.last_epoch + 1
上記の実装はepoch単位で学習率を調整することが前提になっている
5 Training (Attention Is All You Need)の式(3)を実装している
学習の初期段階では0に近い係数(scale)が学習率に掛けられるが、徐々に増えていき、warmup_epochsまでエポック数が達すると係数は1で最大となる。 その後は係数が1より小さくなるため、学習率が減衰されていく。