PyTorch1.4での新機能 Schedulerのchainingを試す

はじめに

みなさんPyTorch1.4にはアップグレードしましたか？まだの方はこちらの公式からアップグレード方法を確認できます。（次のバージョンからPython2系がサポートされなくなるらしいので、注意してくださいね）

PyTorch1.4の新機能としてSchedulerのchaining機能というのがひっそりと追加されていました。（リリースノートはこちら）
早速試してみます。

Schedulerとは

Schedulerを使うと、学習率をEpoch毎に変化させることができます。
学習率は高くした方が早く学習が進むのですが、学習率が高すぎるままだと、最適解を飛び越してしまう恐れがあります。なのでNNの学習時にはSchedulerを使い、Epoch数が進むにつれて徐々に学習率を下げていくのが定石になっています。
（今回の話とは直接関係ありませんが、PyTorchのスケジューラはKeras等と違って、元々の学習率に対する倍率を返すように実装されていることに注意してください）

新機能 Chainingとは

公式によると、「2つのスケジューラーを次々に定義およびステップ実行して、効果を複合できる機能」とのこと。今までは1つのスケジューラが決めた学習率しか使えなかったのに対して、合わせ技ができるということらしいです。
いまいちピンとこないので実際に動かして学習率の変化の様子を見てみましょう。

まずはPyTorch1.3だとどういう挙動だったかを確認

import torch
from torch.optim import SGD
from torch.optim.lr_scheduler import ExponentialLR, StepLR
model = [torch.nn.Parameter(torch.randn(2, 2, requires_grad=True))]
optimizer = SGD(model, 0.1)
scheduler1 = StepLR(optimizer, step_size=3, gamma=0.5)
scheduler2 = ExponentialLR(optimizer, gamma=0.9)
s1, s2, lr = [], [], []
for epoch in range(100):
    optimizer.step()
    scheduler1.step()
    scheduler2.step()
    s1.append(scheduler1.get_lr()[0])
    s2.append(scheduler2.get_lr()[0])
    for param_group in optimizer.param_groups:
        lr.append(param_group['lr'])

StepLRとExponentialLRというスケジューラを２つ使います。それぞれscheduler1,2とします。
得られたスケジューラの学習率(s1, s2)をそれぞれプロットします。

import matplotlib.pyplot as plt
import seaborn as sns
sns.set()
plt.plot(s1, label='StepLR (scheduler1)')
plt.plot(s2, label='ExponentialLR (scheduler2)')
plt.legend()

お互いのスケジューラがどういう特性をもっているかがわかりますね。
次にOptimizerの学習率をプロットします。

plt.plot(lr, label='Learning Rate')
plt.legend()

見ると一目瞭然ですが、ExponentialLR(scheduler2)の学習率が使われているのがわかります。
どうやらPyTorch1.3では最後にstepが呼び出されたschedulerの学習率を使用していたようです。
また、お互いのスケジューラも特に影響しあわず独立に動いていたようです。

いよいよPyTorch1.4でChainingの効果を確認

import torch
from torch.optim import SGD
from torch.optim.lr_scheduler import ExponentialLR, StepLR
model = [torch.nn.Parameter(torch.randn(2, 2, requires_grad=True))]
optimizer = SGD(model, 0.1)
scheduler1 = StepLR(optimizer, step_size=3, gamma=0.5)
scheduler2 = ExponentialLR(optimizer, gamma=0.9)
s1, s2, lr = [], [], []
for epoch in range(100):
    optimizer.step()
    scheduler1.step()
    scheduler2.step()
    s1.append(scheduler1.get_last_lr()[0])
    s2.append(scheduler2.get_last_lr()[0])
    for param_group in optimizer.param_groups:
        lr.append(param_group['lr'])

PyTorch1.3では.get_lr()でスケジューラの学習率を得ていたのですが、PyTorch1.4では.get_last_lr()になっていることに注意してください。
.get_lr()も使えますが、正しい値が出力されないことがあるので注意です。
これらの変更は公式でアナウンスされています。

さて、それぞれのスケジューラの学習率をプロットしてみます。

import matplotlib.pyplot as plt
import seaborn as sns
sns.set()
plt.plot(s1, label='StepLR (scheduler1)')
plt.plot(s2, label='ExponentialLR (scheduler2)')
plt.legend()

この時点で1.3とは様子が違いますね。
続いてOptimizerの学習率をプロットします。

plt.plot(lr, label='Learning Rate')
plt.legend()

このように、２つのスケジューラの学習率が次々と掛け合わさって、最終的なOptimizerの学習率となっていることがわかります。２つのスケジューラ自身もそれぞれ影響しあって学習率が変化している、というのも認識しておくとよさそうですね。各スケジューラの学習率が変化したらOptimizerの学習率も変化するようです。

今まででは自分で書くしかなかった、複雑な学習率の変化をするスケジューラが、簡単に書けるようになりそうです。特に少しCyclicalな挙動を入れたいときなどに重宝しそう。

以上です。