Pytorch学習スケジューラはtimmのCosineLRSchedulerがいいぞ

kaggle的備忘録

Pytorchモデル学習では学習率スケジューラ（LR scheduler）が必須。

特に論文ではCosineLRスケジューラがよく使われる。

こんなん

事前学習モデルでは最初に低いLRでwarmupを数epochした後に最高LRに設定後cosine上にLRが低下していくスケジューリングが経験上良い精度を達成できる。

これがpytorchデフォルトにあれば良いのだが、用意されていない。

一方でtimmに用意されているのは意外と？知られていない。

timm.fast.ai

from timm.scheduler import CosineLRScheduler
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
scheduler = CosineLRScheduler(optimizer, t_initial=100, lr_min=1e-6, 
                                  warmup_t=3, warmup_lr_init=1e-6, warmup_prefix=True)

そして学習ループ内で

for epoch in range(100):
   scheduler.step(epoch)
   train()
   eval()

のように回せば良い。