[Pytorch] Learning Rate Scheduler 사용하기

Contents

 

학습률(Learning Rate)은 모델을 학습하기 위해서 필수적인 요소다. 학습률을 너무 크게 설정한다면, 최솟값에 도달하는 것이 어려우며, 너무 작게 설정하면, local minimum에 빠지거나 학습에 진전이 없을 수 있다. 이번 글에서는 학습률에 Schedular를 설정해서 학습률을 감쇠(Decay)하는 패키지를 다루어볼 것이다. 

 

Import Packages

import torch 
import torch.nn as nn

from torch.optim.lr_scheduler import StepLR 
import torch.optim as optim 

from torch.utils.data import DataLoader  
import torchvision.transforms as transforms 
import torchvision.datasets as dsets

pytorchvision에서 제공하는 datasets 패키지에서 간단한 예제를 불러와 실험할 예정이다. pytorch는 torch.optim. lr_scheduler에서 StepLR 뿐만 아니라 다양한 scheduler를 제공하고 있다. 그 중 가장 대표적인 StepLR의 사용법에 대해서 먼저 알아보자. 

 

# Set seed 
torch.manual_seed(0)
train_dataset = dsets.MNIST(root='./data', 
                            train=True, 
                            transform=transforms.ToTensor(), 
                            download=True)

test_dataset = dsets.MNIST(root='./data', 
                           train=True, 
                           transform=transforms.ToTensor(),     
                           download=True)


batch_size = 100
n_iters = 3000 
num_epochs = n_iters / (len(train_dataset)/batch_size)
num_epochs = int(num_epochs)

train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True)

 

Build Model

class FFNN(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super(FFNN, self).__init__()
        
        self.input_dim = input_dim 
        self.hidden_dim = hidden_dim 
        self.output_dim = output_dim 
        
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.relu = nn.ReLU()
        
        self.fc2 = nn.Linear(hidden_dim, output_dim)
        
        self._init_weight()
        
    def forward(self,images):
        out = self.fc1(images)
        out = self.relu(out)
        out = self.fc2(out)
        return out
    
    def _init_weight(self):
        for module in self.modules():
            if isinstance(module, nn.Linear):
                nn.init.kaiming_normal_(module.weight)
                nn.init.zeros_(module.bias)

모델을 구축하는 것이 핵심이 아니기 때문에 간단한 Feed Forward Neural Network(FFNN)를 구축하였다. initialization은 kaming_he를 사용하였으며, TwoLayer 모델이다.

 

Training

input_dim = 28 * 28
hidden_dim = 100 
output_dim = 10 
learning_rate = 1e-1

model = FFNN(input_dim, hidden_dim, output_dim)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=0.9, nesterov=True)

# step_size를 지정하면 step 마다 learning rate를 gamma만큼 감소시킵니다.
# gamma = decaying factor 
scheduler = StepLR(optimizer, step_size=1, gamma=0.1)

optimizer는 가장 기본적인 Stochastic Gradient Descent(SGD)를 사용하였으며, momentum은 0.9로 설정하였다. StepLR에서 step_size는 몇 번의 step마다 학습률을 조정할 것인지를 의미한다. gamma는 decaying factor로 기존의 학습률에 곱해지는 인수를 의미한다. 

\[ \text{Momentum}: v_t = \gamma u_{t-1} + \eta \cdot \nabla J (\theta - \gamma v_{t-1}, x^{i:i+n}, y^{i:i+n} \]

\[ \theta = \theta - v_t \]

$n$은 batch_size를 의미한다. $\eta$는 학습률을 의미한다. 위 수식을 통해 최적화하고 있으며, 이를 일반화하면 아래와 같다. 

\[ \begin{equation} \begin{split} \text{Given  } \eta_t & = 0.1, \text{ and } \gamma = 0.01 \\ \eta_t & = 0.01 \\ \eta_{t + 1} & = 0.1(0.1) = 0.01 \\ \eta_{t+2} & = 0.1(0.1)^2 = 0.001 \\ \eta_{t + n} & = 0.1(0.1)^n \end{split} \end{equation} \]

 

iter = 0 
for epoch in range(num_epochs):

    print(f'Epoch: {epoch}, LR: {scheduler.get_last_lr()}')
    
    model.train()
    for i, (images, labels) in enumerate(train_loader):
        images = images.view(-1, 28*28)
        
        optimizer.zero_grad()
        pred_y = model(images) # model.train()을 사용하지 않으면, images = imgaes.require_grad_()를 설정해주어야 합니다.
        loss = criterion(pred_y, labels) 
        loss.backward()
        optimizer.step()
        
        iter += 1 
        
        if iter % 500 == 0 :
            correct = 0
            total = 0 
            model.eval()
            for images, labels in test_loader:
                images = images.view(-1, 28*28)
                
                pred_y = model(images)
                pred_y = torch.argmax(pred_y.data, 1)
                
                total += labels.size(0)
                
                correct += (pred_y == labels).sum()
                
            accuracy = 100 * correct / total 
        
            print(f'Iteration: {iter}, Loss: {loss.item():.4f}, Accuracy: {accuracy:.2f}%')
    scheduler.step()
    
    
# Epoch: 0, LR: [0.1]
# Iteration: 500, Loss: 0.2473, Accuracy: 96.47%

# Epoch: 1, LR: [0.010000000000000002]
# Iteration: 1000, Loss: 0.0970, Accuracy: 97.59%

# Epoch: 2, LR: [0.0010000000000000002]
# Iteration: 1500, Loss: 0.0247, Accuracy: 97.75%

# 1Epoch: 3, LR: [0.00010000000000000003]
# Iteration: 2000, Loss: 0.0711, Accuracy: 97.82%

# Epoch: 4, LR: [1.0000000000000004e-05]
# Iteration: 2500, Loss: 0.0305, Accuracy: 97.82%
# Iteration: 3000, Loss: 0.1096, Accuracy: 97.82%

학습률이 점점 감소하는 것을 확인할 수 있다. 사용하는 모델에 따라 gamma를 달리 설정해서 사용하면 된다. torch 버전이 최신 버전인 경우 optimizer.step() 다음에 scheduler.step()을 사용하여야 에러가 발생하지 않으며, get_lr() 대신 get_last_lr()을 사용하여야 한다. 

 

자세한 코드는 [여기]를 참고하길 바란다.