pytorch-lightningの学習サイクルで，WarmupCosineLRSchedulerを使った学習率設定を紹介．
ここでは，エポック単位ではなくてステップ単位の更新を説明します．

準備

timmのインストール

timmのCosineLRSchedulerを使います．
pytorch image modelsという画像認識用のライブラリです

pip install timm

timmのWarmup付き，CosineLRSchedulerクラス自体の使い方は次のサイトがわかりやすいです．

※lightning-boltパッケージにも，WarmupCosineLRReducerがありますがアップデートが全くされておらず，最新のlightningに対応していないので非推奨です．

実装: pl.LightningModuleで使用

Scheduler インスタンスの生成とステップ単位の更新

pl.configure_optimizersメソッドでschedulerのインスタンスを生成して，pl.lr_scheduler_stepでstepメソッドを呼び出すようにしたら完了です．

from timm.scheduler import CosineLRScheduler

class MyLightningModule(pl.LightningModule):
    def __init__(self, model, num_epoch, exp_name) -> None:
        ...
        
        num_dataloader = len(dataloader) # 1epochあたりのバッチ数
        self.total_steps = num_epoch * num_dataloader # 合計step数
        self.current_step = 0

    def configure_optimizers(self):
        optimizer = torch.optim.AdamW(
            self.models, base_lr=1e-4
        )

        warmup_rate = 0.1
        warmup_t = int(self.total_steps * warmup_rate) # warmup step数
        scheduler = CosineLRScheduler(
            optimizer, t_initial=self.total_steps, warmup_t=warmup_t, warmup_prefix=True
        )
        return {
            "optimizer": optimizer,
            "lr_scheduler": {
                "scheduler": scheduler,
                "interval": "step",
                "frequency": 1,
            }
        }

    def lr_scheduler_step(self, scheduler, metric) -> None:
        self.current_step += 1
        scheduler.step(self.current_step)

    ...

設定しているCosineLRSchedulerの引数

t_initial: 総ステップ数　（本来の意味は最初のサイクルのステップ数）
warmup_t: warmupに使うstep数
warmup_prefix: Trueなら，warmup終了時の学習率がOptimizerに設定した値になります

詳細はこちらから

また，返り値のdictではschedulerの更新タイミングや頻度を設定することができます．

"lr_scheduler": {
    "scheduler": scheduler,
    "interval": "step",
    "frequency": 1,
}

今回はstep単位で学習率を更新したいため，"interval": "step"にします．デフォルトは"epoch"なので，この指定は必ず必要です．

最後に，stepメソッドの呼び出しは，LightningModule.lr_scheduler_stepをオーバーライドして実装します．（PytorchのLRSchedulerかlightning-bolt以外のSchedulerを使う場合はオーバーライドしてstepの処理を書く必要があります．）
この関数はschedulerの更新のタイミングで呼ばれるようになっています．

関連

schedulerについてはこちら

LightningModuleとは

LighgningModuleでは， trainステップ， valステップを作成したり，学習におけるサイクルのhooksに任意の処理を追加するAPIが提供されています．
基本的な使い方はこちらの記事をどうぞ