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)](https://github.com/rwightman/pytorch-image-models/blob/master/timm/scheduler/step_lr.py)
>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)
```
![StepLRScheduler.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2473523/0ba25892-b3a4-375c-57e0-844858dabd86.png)

## 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,
)
```
![StepLRScheduler+warmup.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2473523/a7aab4f4-83f0-a557-d35d-423db67e7056.png)

## 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)
```
![StepLRScheduler+noise.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2473523/d2cff98b-39ac-a246-547b-1baee723f7a5.png)


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

Class MultiStepLRScheduler [(source)](https://github.com/rwightman/pytorch-image-models/blob/master/timm/scheduler/multistep_lr.py)
> 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)
```
![MultiStepLRScheduler.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2473523/3d68670c-1789-70c8-e17f-3c788a93836c.png)

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

Class PlateauLRScheduler [(source)](https://github.com/rwightman/pytorch-image-models/blob/master/timm/scheduler/plateau_lr.py)
> 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を使用して学習した結果は以下のようになります。
![plateaulr_loss.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2473523/15af5056-d64c-41c9-7f79-5303be4e9b14.png)
![PlateauLRScheduler.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2473523/d58008cb-fdf1-4853-f4fa-cd400c47a5af.png)

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

Class CosineLRScheduler[(source)](https://github.com/rwightman/pytorch-image-models/blob/master/timm/scheduler/cosine_lr.py)
> 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)
```
![CosineLRScheduler.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2473523/b7b72a77-02ad-e72d-ead5-14a898f21e8c.png)

## 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)
```
![CosineLRScheduler+cycle.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2473523/7b4ea723-852f-6340-1364-1823941ac854.png)
t_initial=epoch数/2, cycle_limit=1とすることで学習率を変更するサイクルを1度だけに制限することもできます。
![CosineLRScheduler+cycle1.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2473523/e11fe829-453b-3eba-f509-c04c1d39d524.png)

## 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)
```
![CosineLRScheduler+cycle_decay.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2473523/4a14ec8d-83e5-1af2-ea44-2073a7202bb3.png)

## 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)
```
![CosineLRScheduler+cycle_mul.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2473523/32229c75-47d5-3c08-03a0-639732c141ec.png)

## k_decay
次にk_decayを追加した場合の出力結果です。k_decayについては[論文](https://arxiv.org/abs/2004.05909)で説明されています。
```
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)
```
![CosineLRScheduler+k_decay.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2473523/3a57236b-454a-9cba-b46f-88b9346975f1.png)

また、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()
```
![k_decay.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2473523/247dc8de-ecab-1121-8d05-6fd928d90193.png)

# TanhLRScheduler
timmのTanhLRSchedulerは以下のような動きをします。
Class TanhLRScheduler[(source)](https://github.com/rwightman/pytorch-image-models/blob/master/timm/scheduler/tanh_lr.py)
> 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)
```
![TanhLRScheduler.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2473523/5cfe9b6d-456c-5f5b-8326-404fb0b59630.png)

## lb, ub
次にlb, ubを変更した場合の出力結果です。lb, ubについては[論文](https://arxiv.org/abs/1806.01593)で説明されています。
```
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)
```
![TanhLRScheduler+lbub.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2473523/92fc8dc8-1bc7-9bf5-7129-4e3bce39abba.png)

また、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()
```
![lbub.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2473523/08ad1517-fc0e-d187-b76f-2c924fd7667d.png)


# 参考文献
[rwightman/pytorch-image-models](https://github.com/rwightman/pytorch-image-models)
[Getting Started with PyTorch Image Models (timm): a practitioner's guide](https://notebooks.githubusercontent.com/view/ipynb?color_mode=auto&commit=49b4ab1af5deec6a90652ac0a9436647ba283b4f&enc_url=68747470733a2f2f7261772e67697468756275736572636f6e74656e742e636f6d2f676973742f76756973656e67392f37626266633335626566636662636261343239636637363138346562353533372f7261772f343962346162316166356465656336613930363532616330613934333636343762613238336234662f74696d6d2e6970796e62&logged_in=false&nwo=vuiseng9%2F7bbfc35befcfbcba429cf76184eb5537&path=timm.ipynb&repository_id=115038248&repository_type=Gist)