PyTorchのSchedulerまとめ

[追記：2019/07/24] 最新版更新してます。

katsura-jp.hatenablog.com

目次

最近暇な時間にPyTorchのReferenceを細かくみたり実装をみたりしているのですが、 今回スケジューラを見ていて、グラフとかあった方嬉しいね(百聞は一見にしかず)と思ったので、グラフをまとめて作りました。

Pytorchのscheduler公式ドキュメントは こちら

PyTorchライブラリ内にあるscheduler

PyTorchでもともと存在するschedulerは以下のとおり

• LambdaLR
• StepLR
• MultiStepLR
• ExponentialLR
• CosineAnnealingLR
• ReduceLROnPlateau

今回は、ReduceLROnPlateau以外の可視化を説明を交えて紹介します。 (ReduceLROnPlateauはloss等が関係するので可視化がめんど...　難しかったので割愛で)

基本設定

今回の記事で使われるパラメータとして

• epoch数: 100
• 初期の学習率: 0.01

model = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
optimizer = optim.SGD(model.parameters(), lr=0.01, weight_decay=0.01, momentum=0.9)

LambdaLR

---

引数一覧

• optimizer : 最適化のインスタンスを指定
• lr_lambda : ラムダ式や関数を指定
• last_epoch : 指定したことないのでいじらなくていいと思う(個人的意見)。深く理解したい人はソースを見てください。

このスケジューラはlr_lambdaにオリジナルの式を入れられるのが最大のポイントです。引数には、stepが呼ばれた回数が与えられ、戻り値としてbase_lr(optimizerに指定したlr)からどれほど変化させたいかの比率を返すようにしてください。つまり、学習率はbase_lr * 戻り値となります。

example

• ラムダ式を与えた場合

scheduler = LambdaLR(optimizer, lr_lambda = lambda epoch: 0.95 ** epoch)
for epoch in range(0, 100): #ここは以下省略
    scheduler.step() 

• 関数を渡した場合

def func(epoch):
    if epoch < 40:
        return 0.5
    elif epoch < 70:
        return 0.5**2
    elif epoch < 90:
        return 0.5**3
    else:
        return 0.5**4
scheduler = LambdaLR(optimizer, lr_lambda = func)

StepLR

---

引数一覧

• optimizer : 省略
• step_size : 何ステップごとに減衰させるかの値
• gamma : 減衰率
• last_epoch : 省略

これは見てもらった方が早いと思います。

example

scheduler = StepLR(optimizer, step_size=20, gamma=0.5)

MultiStepLR

---

引数一覧

• optimizer : 省略
• milestones : 減衰させたいstepのリスト
• gamma : 減衰率
• last_epoch : 省略

StepLRは減衰ステップが一つに対し、これは複数取れます。注意点として、milestonesには、ステップの小さい順のリストを与えてください。

つまり、10,30,50のステップ数で減衰させたい場合は、[10,30,50]と与えてください。

example

scheduler = MultiStepLR(optimizer, milestones=[10, 20, 40], gamma=0.5) # 10 < 20 < 40

ExponentialLR

---

引数一覧

• optimizer : 省略
• gamma : 減衰率
• last_epoch : 省略

簡潔に説明するとbase_lr**step回数が学習率に与えられる。

example

scheduler = ExponentialLR(optimizer, gamma=0.95)

CosineAnnealingLR

---

引数一覧

• optimizer : 省略
• T_max : 1 半周期のステップサイズ
• eta_min : 最小学習率(極小値)
• last_epoch : 省略

これは、kaggleやsignateの画像コンペでしばしば解法で登場するsnapshot ensemblingで使用されてます。もちろんそのほかでも使われます。

秋頃(2018年9月ころ)にpytorchになかった気もしなくもなく(見つけてなかっただけかも)、じぶんで実装した記憶があります。

これも見ていただいた方がわかりやすい

example

scheduler = CosineAnnealingLR(optimizer, T_max=20, eta_min=0.001)

ReduceLROnPlateau

---

引数一覧

• optimizer : 省略
• mode : 監視されている要因が'min'だと下がっているか、'max'だと上がっているかを指定
• factor : 学習率の減衰率
• patience : 何ステップ向上しなければ減衰するかの値
• verbose : 減衰するとき出力してれるかどうか
• threshold : threshold_modeで説明するのに使われてる値
• threshold_mode : 'rel'と'abs'の2種類のどちらかを取ります。 'rel'の場合、減衰の起因となる閾値をmaxモードだとbest*(1+threshold)、minモードだとbest*(1-threshold)となります。'abs'の場合は、maxモードだとbest-threshold、minモードだとthresh+thresholdとなります。
• cooldown : 学習率が下がってから、factorを再監視するまでのステップ数
• min_lr : 最小学習率
• eps : nanとかInf回避用の微小数

また、監視する要因はschedulerの関数に存在するstepの引数に与えます。

また、大げさに例をあげると、 - mode : 'min' - patience : 10 - threshold : 0.5 - threshold_mode : 'abs' のように指定すると、監視している値のbestが0.01とすると、 0.01+0.5=0.51よりも高い値を10回連続で取ると減衰します。

example

scheduler = ReduceLROnPlateau(optimizer, 'min') 
for epoch in range(10):
    scheduler.step(val_loss) #val_lossが下がらなければ減衰

自作scheduler

---

PyTorchでは自分でschedulerを作成することも容易です。torch.optim.lr_scheduler._LRSchedulerを継承したクラスを作成すると、上に紹介したようなschedulerを自作することが容易にできます。

今回わかりやすい例として、Loss Surfaces, Mode Connectivity, and Fast Ensembling of DNNs で使われているようなlinear cyclical learning rate scheduleを作成して見たいと思います。

実装

from torch.optim.lr_scheduler import _LRScheduler

class LinearCyclicalLR(_LRScheduler):
        def __init__(self, optimizer, T_max, eta_min=0, last_epoch=-1):
            self.T_max = T_max
            self.eta_min = eta_min
            super(LinearCyclicalLR, self).__init__(optimizer, last_epoch)

        def get_lr(self):
            last_step = self.last_epoch % self.T_max
            point = self.T_max // 2
            return [-(base_lr-self.eta_min)* last_step / (point-0.5) + base_lr if last_step < point
                    else (base_lr-self.eta_min)*last_step/(point-0.5) + self.eta_min - base_lr
                    for base_lr in self.base_lrs]

まず、pytorchのschedulerの自作にあたり、_LRSchedulerを継承します。optimizerの学習率更新は、scheduler.step()が呼ばれた際に更新されますが、その際、self.last_epochがインクリメントされます。厳密には、このget_lrが呼ばれる前に行われます。step内でself.get_lrが呼ばれ、その戻り値を新たな学習率に置き換えるように設計されてます。

つまり、自作関数でschedulerを作成するならば、get_lrを自作するのにほとんど等しいです。

今回のget_lrの戻り値のリスト内包表記は中学1年で習うような一次関数なので説明は省略。引数はCosineAnnealingLRを丸パクリ。

これを実際に学習率をプロットすると

scheduler = LinearCyclicalLR(optimizer, T_max=20, eta_min=0.001)

おわりに

---

いかがだったでしょうか。 最近TLにPyTorchはじめた方が散見されていましたので書いて見ました。個人的にはPyTorch書きやすい/読みやすい/普及しているの3点でおすすめです。 ちなみに、学習率の取得は、scheduler.get_lr()で返ってきます。

しかし、chainer(chainercvなども含める)と比べられると自分で実装しないといけないことが多いかもしれません。どちらも使えるように自分も頑張りたいと思います。