[PyTorch] CosineAnnealingLR, CosineAnnealingWarmRestarts

Cosine Annealing은 SGDR: Stochastic Gardient Descent with Warm Restarts에서 제안된 learning rate scheduling 방법이다.

 

Cosine Annealing은 learning rate가 cosine 함수의 주기를 따라 감소하고, 증가하는 과정을 반복한다.

 

우선, Cosine Annealing의 수식은 다음과 같다.

 

$\eta_t=\eta_{min}+\frac12(\eta_{max}-\eta_{min})(1+cos(\frac{T_{cur}}{T_{max}}\pi))$

 

$\eta$는 learning rate를 나타낸다.

 

Cosine Annealing에서 설정 가능한 parameter은 다음과 같다.

1. $\eta_{min}$ : 최소 learning rate 값

2. $\eta_{max}$ : 최대 learning rate 값

3. $T_{max}$ : 한 주기의 길이

 

위 그림에서는 $\eta_{min}$=0, $\eta_{max}$=0.1, $T_{max}$=100로 설정되어 있다.

 

 

CosineAnnealingLR

PyTorch에서는 torch.optim.lr_scheduler.CosineAnnealingLR을 이용해 다음과 같이 사용할 수 있다.

 

 

 

T_max와 eta_min은 위의 수식과 동일하다.

eta_max는 scheduler에서 따로 설정하지 않고, optimizer에 설정된 learing rate를 사용한다.

import torch.optim as optim

epochs = 300
optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100, eta_min=0)
for epoch in range(epochs):
    train(...)
    val_loss = validate(...)
    scheduler.step()

 

 

CosineAnnealingWarmRestarts

torch.optim.lr_scheduler.CosineAnnealingWarmRestarts를 이용하면 Warm Restart를 추가할 수 있다. Warm Restart는 learning rate의 감소 후 급격하게 증가하게 함으로써 local optima에 빠지는 것과 overfitting을 방지할 수 있다.

 

 

 

CosineAnnealingLR의 T_max 대신 T_0와 T_mult가 있고, eta_min은 동일한데,

 

T_0은 첫 주기의 길이를 의미하고, 이후에는 주기가 T_mult만큼 곱해진다. 단 T_mult는 integer만 가능하여 주기가 짧아지는 것은 불가능하다.

import torch.optim as optim

epochs = 300
optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=50, T_mult=2, eta_min=0)
for epoch in range(epochs):
    train(...)
    val_loss = validate(...)
    scheduler.step()