[DL] PyTorch에서 Learning Rate Scheduler 그리기

Deep learning^딥러닝 은 기본적으로 gradient descent^{경사하강법} 기반의 알고리즘을 이용하여 Figure 1과 같이 모델의 weight을 최적화한다. 이 과정에서 적당한 learning rate (LR) ^학습률을 설정하여 학습하는 것이 중요한데, 너무 적게 설정하면 모델이 수렴하는 데에 시간이 너무 오래걸리고, 너무 크게 학습하면 global minima를 향해 가지 못하여 수렴하지 못할 수도 있다.

Figure 1 LR에 따른 모델 수렴
(출처 : bdhammel.com)

Learning rate을 모델을 학습하는 동안 처음부터 끝까지 동일하게 설정해줄 수도 있지만, 최적값에 가까워질수록 learning rate을 줄여가며 학습하는 등 lr을 바꿔가며 학습하는 방법도 있다. 다행히도 그동안 learning rate을 스케줄링해주는 learning rate scheduler에 대한 다양한 연구들이 많이 진행되어 왔고, PyTorch 공식 framework에 torch.optim.lr_scheduler (link)에 구현이 되어있다. 하지만 이 코드들이 대부분 잘 구현이 되어있긴 하지만, 내 입맛에 맞게 customizing해야 하는 경우도 있다. 여기서는 이 learning rate scheduler들에 대한 자세한 설명보다는, 쓰고자 하는 learning rate schduler를 사전에 plot해보면서 눈으로 확인해보는 방법에 대해 알아보려 한다. 학습 전에 내가 설정한 lr scheduler를 미리 보는 방법을 몰라서 삽질했던 경험이…;

PyTorch Framework에 구현되어 있는 lr_scheduler 그리기

PyTorch에 이미 구현되어 있는 lr scheduler를 그려보자!

Cosine Annealing Warm Restarts

연구하면서 AdamW optimizer와 Cosine Annealing Warm Restarts lr scheduler를 이용하여 학습할 때가 가장 안정적으로 학습되었던 경험이 있었다. Cosine Annealing Warm Restarts은 아래 그림 Figure 2와 같이 다시 lr을 크게 해줌으로써 고질적인 local minima에서 탈출시켜 global minima를 향해 수렴할 수 있기 때문으로 보인다. 따라서, 이를 위주로 설명해보고자 한다.

Figure 2 PyTorch Cosine Annealing Warm Restarts

(1) 필요한 Library Import

import torch
import torch.nn as nn
from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts
import matplotlib.pyplot as plt

(2) Model, Optimizer선언

1. 아주 간단한 모델 하나를 선언해주자. (아무 모델이나 상관 없음, 여기서는 Fully-connected layer를 선언해줌)

   model = nn.Linear(2, 1)
   

2. Optimizer도 아무거나 선언해주자!

   optimizer = torch.optim.AdamW(model.parameters())
   

(3) LR Scheduler 선언

scheduler = CosineAnnealingWarmRestarts(optimizer,\
			T_0=30, T_mult=1, eta_min=1e-6, last_epoch=-1)

먼저 앞서 선언한 optimizer 객체를 인자로 넣어준다.

• T_0=30 : Cosine annealing의 첫 번째 iteration의 주기가 30
• T_mult=1 : 각 iteration이 재시작 될 때마다 주기를 몇 배수로 늘릴지 -> 1배수로 주기 모두 동일
• eta_min=1e-6 : Annealing 시 가질 수 있는 learning rate 최솟값
• last_epoch=-1 : 얜 그냥 건들지 말자

(3) Step을 나아가며 저장 후 print

lrs  = []
for  i  in  range(149):
	scheduler.step()
	lrs.append(optimizer.param_groups[0]["lr"])  

plt.plot(lrs)
plt.xlabel("Epochs")
plt.ylabel("Learning Rate")
plt.title("CosineAnnealingWarmRestarts LR Scheduler")

사실 공식 코드는 아래 그림 Figure 3와 같이 실제 논문 SGDR에서 제안한 것과 동일하게 구현되어 있지 않다. 공식 코드에서 튜닝할 수 없는 부분은 다음과 같다.

Figure 3 Customized Cosine Annealing Warm Restarts

1. 각 iteration의 max learning rate (Eta_max)을 감쇄시킬 수 없다.
2. 각 iteration의 주기를 늘릴 때 정수배로 밖에 설정해줄 수 없다.
3. Literally warm restart 부분이 구현되어 있지 않다. Warm restarts를 사용하면 Transformer 계열 모델을 학습시킬 때 안정적으로 학습시킬 수 있다고 한다.

해결책은 katsura-jp의 repo의 lr scheduler를 이용하면 이 모든 사항을 반영할 수 있다.