PyTorch Scheduler Class の Chaining

今回はPyTorch 1.4.0から追加された新機能であるSchedulerのChainingについてです。この方のQiita記事で知ったのですが、PyTorchのschedulerに新機能が追加された様です。

qiita.com

内容はシンプルで、今までは学習ループの中に1種類しか組み込めなかったschedulerを複数重ねて使えると言うことです。具体例をみていきましょう。

import torch
from torch.optim.lr_scheduler import ExponentialLR, CosineAnnealingLR
from torch.optim import SGD
import matplotlib.pyplot as plt
import seaborn as sns
sns.set()

lr_cos, lr_exp, lr = [], [], []

# 初期化
model = [torch.nn.Parameter(torch.randn(2, 2, requires_grad=True))]
optimizer = SGD(model, 1e-2)
optimizer.zero_grad()
scheduler_cos = CosineAnnealingLR(optimizer, T_max=20, eta_min=1e-6)
for epoch in range(200):
    optimizer.step()
    optimizer.zero_grad()
    scheduler_cos.step()
    lr_cos.append(scheduler_cos.get_last_lr()[0])

# 初期化
model = [torch.nn.Parameter(torch.randn(2, 2, requires_grad=True))]
optimizer = SGD(model, 1e-2)
optimizer.zero_grad()
scheduler_exp = ExponentialLR(optimizer, gamma=0.97)
for epoch in range(200):
    optimizer.step()
    optimizer.zero_grad()
    scheduler_exp.step()
    lr_exp.append(scheduler_exp.get_last_lr()[0])

# 初期化
model = [torch.nn.Parameter(torch.randn(2, 2, requires_grad=True))]
optimizer = SGD(model, 1e-2)
optimizer.zero_grad()
scheduler_cos = CosineAnnealingLR(optimizer, T_max=20, eta_min=1e-6)
scheduler_exp = ExponentialLR(optimizer, gamma=0.97)
for epoch in range(200):
    optimizer.step()
    optimizer.zero_grad()
    scheduler_exp.step()
    scheduler_cos.step()
    lr.append(optimizer.param_groups[0]['lr'])

plt.plot(lr_cos, label='CosineAnnealing', linestyle='dashed')
plt.plot(lr_exp, label='Exponential', linestyle='dashed')
plt.plot(lr, label='Fusion', color='red', linewidth=3)
plt.legend(bbox_to_anchor=(0, -0.1), loc='upper left', borderaxespad=0)
plt.show()

この様な形でCosineAnnealingとExponentialを使って新たなタイプのschedulerを作れていることが分かります。この機能と以下の方のコードを使えば、今まで自分でクラスを作っていたWarmup系のschedulerも簡単に作れそうです。

github.com

!pip install git+https://github.com/ildoonet/pytorch-gradual-warmup-lr.git
from warmup_scheduler import GradualWarmupScheduler

lr = []

# 初期化
model = [torch.nn.Parameter(torch.randn(2, 2, requires_grad=True))]
optimizer = SGD(model, 1e-2)
optimizer.zero_grad()
scheduler_cos = CosineAnnealingLR(optimizer, T_max=400, eta_min=1e-6)
scheduler_wup = GradualWarmupScheduler(optimizer, multiplier=1, total_epoch=20, after_scheduler=scheduler_cos)
for epoch in range(200):
    optimizer.step()
    optimizer.zero_grad()
    scheduler_cos.step()
    scheduler_wup.step()
    lr.append(optimizer.param_groups[0]['lr'])

plt.plot(lr, label='Warmup Cosine Annealing', color='red', linewidth=3)
plt.legend(bbox_to_anchor=(0, -0.1), loc='upper left', borderaxespad=0)
plt.show()