Pytorchのtimm.schedulerを試してみた

準備

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

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