[PyTroch] Optimizing Model Parameters

이번 글에서는 data를 사용해 모델 parameter들을 최적화 하는 방법을 소개한다. 모델을 학습은 반복 과정이다. 각 iteration에서 모델은 output의 예측값을 만들고, 에러를 계산하고 (loss), 각 parameter에 대해 미분값을 계산하고, parameter들을 gradient descent 방법을 사용해 최적화한다.

 

Prerequisite Code

Datasets & DataLoaders와 Build Model 섹션의 코드이다.

import torch
form torch import nn
from troch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transform import ToTensor

training_data = dataset.FashionMNIST(
	root="data",
    train=True,
    download=True,
    transform=ToTensor()
)

test_data = datasets.FashionMNIST(
	root="data",
    train=False,
    download=True,
    transform=ToTensor()
)

train_dataloader = DataLoader(training_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)

class NeuralNetwork(nn.Module):
	def __init__(self):
    	super().__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
        	nn.Linear(28*28, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 10),
        )
    
    def forward(self, x):
    	x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits
    
model = NeuralNetwork()

 

Hyperparameters

Hyperparameter들은 모델 최적화 과정을 조절할 수 있게 해주는 parameter들이다. 다른 hyperparameter 값들은 모델 학습과 수렴 속도에 영향을 줄 수 있다.

 

학습에는 다음과 같은 hyperparameter들을 정의한다.

• Number of Epochs - 전체 데이터셋을 순회하는 횟수
• Batch Size - parameter 업데이트 전에 네트워크에 전파되는 데이터 샘플 수
• Learning Rate - 각 batch/epoch마다 모델 parameter를 얼마나 업데이트 할지

learning_rate = 1e-3
batch_size = 64
epochs = 5

 

Optimization Loop

Hyperparameter들을 정한 후 optimization loop를 통해 모델을 최적화할 수 있다. Optimization loop의 각 iteration은 epoch라고 불린다.

 

각 epoch는 두 가지 주요 부분으로 구성되어있다.

• Train Loop - 학습 데이터셋을 순회하면서 최적화된 parameter들에 수렴을 목표로 함
• Validation/Test Loop - 테스트 데이터셋을 순회하면서 모델의 성능이 향상되고 있는지 확인을 목표로 함

 

이 학습 loop에 사용되는 몇가지 개념들을 간단하게 알아보자.

Loss Function

Loss function은 모델이 출력한 결과와 목표 값의 다른 정도를 측정한다. 즉, 학습 과정에서는 이 loss function 최소화를 목표로 한다. Loss를 계산하기 위해서 주어진 input 데이터 샘플들로 예측값을 내고 실제 데이터 label과 비교하게 된다.

 

주로 사용되는 loss function들은 다음과 같다.

• nn.MSELoss (Mean Squared Error) - regression 작업
• nn.NLLLoss (Negative Log Likelihood) - classification 작업
• nn.CrossEntropyLoss - nn.LogSoftmax와 nn.NLLLoss를 결합한 loss

 

예제에서는 output logit들을 nn.CrossEntropyLoss로 보낸다. 이는 logit들을 정규화사키고 예측값의 에러를 계산한다.

# loss 함수 초기화
loss_fn = nn.CrossEntropyLoss()

 

Optimizer

최적화는 각 학습 단계에서 모델의 parameter들을 조정해 모델 에러를 줄이는 과정이다. 최적화 알고리즘은 이 과정이 어떻게 실행되는지 정의한다. (아래 예제에서는 Stochastic Gradient Descent를 사용한다.) 모든 최적화 로직은 optimizer 객체에 포함되어있다. PyTorch에는 SGD optimizer 뿐만 아니라 ADAM과 RMSProp 등 서로 다른 모델 데이터에 더 잘 작동하는 optimzer들을 가지고 있다.

 

Optimizer는 학습되어야하는 parameter들과 learning rate hyperparameter를 전달해 초기화한다.

optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

 

학습 loop 안에서 최적화는 세 가지 단계에 걸쳐 이루어진다.

• optimizer.zero_grad()를 호출해 모델 parameter들의 기울기를 초기화한다. 기울기들은 기본적으로 모두 더해지기 때문에 중복을 피하려면 각 iteration마다 0으로 초기화해야 한다.
• loss.backward()를 호출해 예측 loss를 backpropagate 한다. PyTorch는 각 parameter에 대한 loss를 저장한다.
• 기울기를 계산한 후 optimizer.step()을 호출해 backward pass에서 수집한 기울기들에 기반해 parameter들을 조절한다.

 

Full Implementation

아래 예제에서는 최적화 코드를 순회할 train_loop와 테스트 데이터로 모델의 성능을 평가할 test_loop를 정의한다.

def train_loop(dataloader, model, loss_fn, optimizer):
	size = len(dataloader.dataset)
    
    # 모델을 training mode로 변경 - batch 정규화와 dropout layer들에 중요
    # 이 상황에서는 필요하지 않지만 모범 사례로 추가됨
    model.train()
    for batch, (X, y) in enumerate(dataloader):
    	# 예측값과 loss를 계산
        pred = model(X)
        loss = loss_fn(pred, y)
        
        # Backpropagation
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
        
        if batch % 100 == 0:
        	loss, current = loss.item(), batch * batch_size + len(X)
            print(f"loss: {loss:>7f}	[{current:>5d|size:>5d}]")
            

def test_loop(dataloader, model, loss_fn):
    # 모델을 evaluation mode로 변경 - batch 정규화와 dropout layer들에 중요
    # 이 상황에서는 필요하지 않지만 모범 사례로 추가됨
    model.eval()
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    test_loss, correct = 0, 0
    
    # 모델을 torch.no_grad()로 실행하면 test mode에서 기울기가 계산되지 않게 함
    # 또한 requires_grad=True로 설정된 tensor들이 사용하는 필요 없는 기울기 계산과 메모리 사용량을 줄임
    with torch.no_grad():
    	for X, y in dataloader:
        	pred = model(X)
            test_loss += loss_fn(pred, y).item()
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()
            
    test_loss /= num_batches
    correct /= size
    print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")

 

다음 예제에서는 loss function과 optimizer를 초기화하고 train_loop와 test_loop에 전달한다.

loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

epochs = 10
for t in range(epochs):
	print(f"Epoch {t+1}\n--------------------------------")
    train_loop(train_dataloader, model, loss_fn, optimizer)
    test_loop(test_dataloader, model, loss_fn)
print("DONE!")