準備

まず今回使用するモジュールをインポートします。

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.optim as optim
import timm
import timm.scheduler

次にshedulerをスムーズに確認するための関数を定義しておきます。

def create_model_and_optimizer():
    model = torch.nn.Linear(2,1)
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-01)
    return model, optimizer

def create_epochs_and_others(
    num_epochs=100, repeat = 1, num_steps_per_epoch = 10,
):
    num_epochs = num_epochs
    num_epoch_repeat = num_epochs//repeat
    num_steps_per_epoch = num_steps_per_epoch
    return num_epochs, num_epoch_repeat, num_steps_per_epoch

def plot_lrs(scheduler, println=False):
    lrs = []
    for epoch in range(num_epochs):
        num_updates = epoch * num_steps_per_epoch
        
        for i in range(num_steps_per_epoch):
            if println:
                print(f"epoch:{epoch:03d}, iter:{i:03d}, lr:{optimizer.param_groups[0]['lr']}")
            num_updates += 1
            scheduler.step_update(num_updates=num_updates)
        scheduler.step(epoch+1)#次のepoch数を示す 
        lrs.append(optimizer.param_groups[0]["lr"])
    plt.plot(lrs)
    return lrs

StepLRScheduler

timmのStepLRSchedulerはtorch.optim.lr_scheduler.StepLRと似た以下のような動きをします．

Class StepLRScheduler (source)

StepLRScheduler(optimizer: torch.optim.Optimizer, decay_t: float, decay_rate: float = 1.,
warmup_t=0, warmup_lr_init=0, t_in_epochs=True,noise_range_t=None,
noise_pct=0.67, noise_std=1.0, noise_seed=42,initialize=True)

デフォルト

StepLRSchedulerをそのまま出力した結果です。decay_t毎にdecay_rateをかけているのが確認できます。

num_epochs, num_epoch_repeat, num_steps_per_epoch = create_epochs_and_others()
model, optimizer = create_model_and_optimizer()

scheduler = timm.scheduler.StepLRScheduler(
    optimizer,
    decay_t =30, # 何epoch毎にlrにdecay_rateをかけるか
    decay_rate = 0.5, # lrの減衰率
    t_in_epochs=True
)
_ = plot_lrs(scheduler)

warmup

次にwarmupを追加した場合の出力結果です。warmup_tで指定したepoch数の間で学習率が上昇しているのが確認できます。

model, optimizer = create_model_and_optimizer()

scheduler = timm.scheduler.StepLRScheduler(
    optimizer,
    decay_t =30,
    decay_rate = 0.5,
    warmup_t = 10, # warmupをかけるepoch数
    warmup_lr_init=1e-04, # warmup開始時のlrの初期値
    t_in_epochs=True,
)

noise

次にnoiseを追加した場合の出力結果です。

model, optimizer = create_model_and_optimizer()

scheduler = timm.scheduler.StepLRScheduler(
    optimizer,
    decay_t = 30,
    decay_rate = 0.5,
    noise_range_t = [0,30], #noiseをかける範囲
    noise_pct = 0.1, #default 0.67
    noise_std = 1.0, #default
    t_in_epochs=True,
)

_ = plot_lrs(scheduler)

MultiStepLRScheduler

timmのMultiStepLRSchedulerはtorch.optim.lr_scheduler.MultiStepLRと似た以下のような動きをします．

Class MultiStepLRScheduler (source)

MultiStepLRScheduler(optimizer: torch.optim.Optimizer,decay_t: List[int],decay_rate: float = 1.,
warmup_t=0,warmup_lr_init=0,t_in_epochs=True, noise_range_t=None,
noise_pct=0.67,noise_std=1.0,noise_seed=42,initialize=True,)

デフォルト

MultiStepLRSchedulerをそのまま出力した結果です。decay_tで指定した値でdecay_rateをかけているのが確認できます。

num_epochs, num_epoch_repeat, num_steps_per_epoch = create_epochs_and_others()
model, optimizer = create_model_and_optimizer()

scheduler = timm.scheduler.MultiStepLRScheduler(
    optimizer,
    decay_t = [20,50,90],
    decay_rate = 0.5,
    t_in_epochs=True,
)

_ = plot_lrs(scheduler)

PlateauLRScheduler

timmのPlateauLRSchedulerはtorch.optim.lr_scheduler.ReduceLROnPlateauと似た以下のような動きをします。
lossを監視して条件を満たせば学習率を係数で減衰させます。

Class PlateauLRScheduler (source)

PlateauLRScheduler(optimizer,decay_rate=0.1,patience_t=10,verbose=True,
threshold=1e-4,cooldown_t=0,warmup_t=0,warmup_lr_init=0,lr_min=0,
mode='max',noise_range_t=None,noise_type='normal',noise_pct=0.67,
noise_std=1.0,noise_seed=None,initialize=True,)

例

実際にPlateauSchedulerを使用して学習した結果は以下のようになります。

CosineLRScheduler

timmのCosineLRSchedulerはtorch.optim.lr_scheduler.CosineAnnealingWarmRestartsと似た以下のような動きをします。

Class CosineLRScheduler(source)

CosineLRScheduler(optimizer: torch.optim.Optimizer,t_initial: int, lr_min: float = 0.,
cycle_mul: float = 1.,cycle_decay: float = 1.,cycle_limit: int = 1, warmup_t=0,
warmup_lr_init=0,warmup_prefix=False,t_in_epochs=True,noise_range_t=None,
noise_pct=0.67,noise_std=1.0,noise_seed=42,k_decay=1.0,initialize=True)

デフォルト

timmのCosineLRSchedulerをそのまま出力した結果です。

num_epochs, num_epoch_repeat, num_steps_per_epoch = create_epochs_and_others()
model, optimizer = create_model_and_optimizer()
scheduler = timm.scheduler.CosineLRScheduler(
    optimizer,
    t_initial=num_epochs,
    lr_min=1e-04,
    t_in_epochs=True,
)
_= plot_lrs(scheduler)

cycle_limit

次にcycle_limitを追加した場合の出力結果です。
t_initial=epoch数/2, cycle_limit=2とすることで学習率を変更するサイクルを２度繰り返しています。

num_epochs, num_epoch_repeat, num_steps_per_epoch = create_epochs_and_others(repeat=2)
model, optimizer = create_model_and_optimizer()
scheduler = timm.scheduler.CosineLRScheduler(
    optimizer,
    t_initial=num_epoch_repeat, # num_epoch_repeat = num_epochs//2
    lr_min=1e-04,
    cycle_limit= 2,
    t_in_epochs=True,
)
_= plot_lrs(scheduler)

t_initial=epoch数/2, cycle_limit=1とすることで学習率を変更するサイクルを1度だけに制限することもできます。

cycle_decay

次にcycle_decayを追加した場合の出力結果です。
学習率を変更するサイクルが繰り返されるたびに学習率にcycle_decayをかけています。

num_epochs, num_epoch_repeat, num_steps_per_epoch = create_epochs_and_others(repeat=3)
model, optimizer = create_model_and_optimizer()
scheduler = timm.scheduler.CosineLRScheduler(
    optimizer,
    t_initial=num_epoch_repeat, # num_epoch_repeat = num_epochs//3
    lr_min=1e-04,
    cycle_limit= 3,
    cycle_decay=0.8,
    t_in_epochs=True,
)
_= plot_lrs(scheduler)

cycle_mul

次にcycle_mulを追加した場合の出力結果です。
学習率を変更するサイクルが繰り返されるたびにt_initialにcycle_mulをかけています。

num_epochs, num_epoch_repeat, num_steps_per_epoch = create_epochs_and_others(repeat=3)
model, optimizer = create_model_and_optimizer()
scheduler = timm.scheduler.CosineLRScheduler(
    optimizer,
    t_initial=num_epoch_repeat, # num_epoch_repeat = num_epochs//3
    lr_min=1e-04,
    cycle_limit= 2,
    cycle_mul=2.0,
    t_in_epochs=True,
)
_= plot_lrs(scheduler)

k_decay

次にk_decayを追加した場合の出力結果です。k_decayについては論文で説明されています。

model, optimizer = create_model_and_optimizer()
scheduler = timm.scheduler.CosineLRScheduler(
    optimizer,
    t_initial=num_epochs,
    lr_min=1e-04,
    t_in_epochs=True,
    k_decay=2.0,
)
_= plot_lrs(scheduler)

また、k_decayを理解するために軽く実装してみました。

import math

T = num_epochs
lr_min = 0.01
lr_max = 0.1
k_decay = 2.0

lr_list = []
for t in range(0, T):
    i = t//T
    t_curr = t - (T*i)
    lr = lr_min + (lr_max - lr_min)/2 * (1 + math.cos(math.pi * t_curr**k_decay/T**k_decay))
    lr_list.append(lr)

plt.plot(np.arange(len(lr_list)), lr_list)
plt.show()

TanhLRScheduler

timmのTanhLRSchedulerは以下のような動きをします。
Class TanhLRScheduler(source)

TanhLRScheduler(optimizer: torch.optim.Optimizer,t_initial: int,
lb: float = -7.,ub: float = 3.,lr_min: float = 0.,cycle_mul: float = 1.,
cycle_decay: float = 1.,cycle_limit: int = 1,warmup_t=0,warmup_lr_init=0,
warmup_prefix=False,t_in_epochs=True,noise_range_t=None,noise_pct=0.67,
noise_std=1.0,noise_seed=42,initialize=True)

デフォルト

timmのTanhLRSchedulerをそのまま出力した結果です。

num_epochs, num_epoch_repeat, num_steps_per_epoch = create_epochs_and_others()
model, optimizer = create_model_and_optimizer()
scheduler = timm.scheduler.TanhLRScheduler(
    optimizer,
    t_initial=num_epochs,
    lr_min=1e-04,
    t_in_epochs=True,
)
_= plot_lrs(scheduler)

lb, ub

次にlb, ubを変更した場合の出力結果です。lb, ubについては論文で説明されています。

model, optimizer = create_model_and_optimizer()
# assert lb < ub
scheduler = timm.scheduler.TanhLRScheduler(
    optimizer,
    t_initial=num_epochs,
    lb = -7,
    ub = 3,
    lr_min=1e-04,
    t_in_epochs=True,
)
_= plot_lrs(scheduler)

また、lbとubを理解するために軽く実装しました。

import math

T = num_epochs
lr_min = 0.01
lr_max = 0.1
U = 3
L = -7

lr_list = []
for t in range(0, T):
    lr_HTD = lr_min + (lr_max - lr_min)/2 * (1-math.tanh(L+(U-L)*(t/T))) #1
#    lr_HTD = lr_min + (lr_max - lr_min)/2 * (1-math.tanh(L*(1-t/T) + U*(t/T))) #1 の式変形
    lr_list.append(lr_HTD)

plt.plot(np.arange(len(lr_list)), lr_list)
plt.show()

参考文献

rwightman/pytorch-image-models
Getting Started with PyTorch Image Models (timm): a practitioner's guide