[pytorch] Tutorial - Optimizing Model Parameters

References

• Official PyTorch Tutorial (link)

---

이번에는 앞선 튜토리얼에서 살펴본 내용들을 통합하여, 네트워크 모델을 학습, 검증, 테스트하는 방법을 살펴보겠습니다. 모델의 학습은 iterative process입니다. 각 반복에서 모델을 output을 계산하고, 실제 결과와의 loss를 통해 error를 계산하고 파라미터(모델의 weights)에 대한 gradient를 계산한 뒤, gradient descent 알고리즘을 사용하여 파라미터를 최적화합니다.

Prerequisite Code

다음 코드를 통해서 FashionMNIST 데이터를 받아오고, 간단한 네트워크를 구성합니다.

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor

training_data = datasets.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(NeuralNetwork, self).__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

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

Learning rate, 배치 사이트, 반복 횟수를 지정합니다.

 

Optimization Loop

텐서플로우의 keras는 모델을 생성하고 fit() API를 사용하면 간단하게 학습할 수 있지만, 파이토치는 학습 과정을 모두 구현해주어야 합니다. 즉, optimization loop를 구현해주어야 하며, 루프를 통해 네트워크 모델을 학습하고 최적화하게 됩니다. 여기서 루프에서 각 반복을 epoch라고 부릅니다.

 

각 epoch에서는 train loop와 validation/test loop로 구성됩니다.

 

Loss Function

Optimization loop를 구성하기 전에 먼저, 사용할 loss function을 정의해야 합니다. 일반적으로 regression task에서는 nn.MSELoss (Mean Square Error)를 사용하고, classification에서는 nn.NLLLoss (Negative Log Likelihood)를 주로 사용합니다. nn.CrossEntropyLoss는 nn.LogSoftmax와 nn.NLLLoss가 결합된 버전입니다.

 

위에서 정의한 모델의 task는 classification이므로 nn.CrossEntropyLoss를 사용하고, 이 loss function에 모델의 output을 전달하면 logits(output)을 normalization하고 prediction error를 계산합니다.

# Initialize the loss function
loss_fn = nn.CrossEntropyLoss()

 

Optimizer

사용할 optimization algorithm도 정의해주어야 합니다. Optimization은 각 학습 단계에서 모델의 error를 줄이기 위해 모델의 파라미터를 조정하는 과정이며, optimization algorithm은 이러한 과정이 수행되는 방법을 정의합니다. 여기서는 SGD(Stochastic Gradient Descent)를 사용하며, ADAM 등 다양한 종류의 optimization을 사용할 수도 있습니다.

 

Optimization을 생성할 때, 최적화할 파라미터와 하이퍼파라미터를 전달해줍니다.

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

Training loop 내에서 최적화는 다음의 세 단계가 수행됩니다.

• optimizer.zero_grad()를 호출하여 모델 파라미터의 gradients를 리셋한다. 기본적으로 gradient는 합산되며, 중복 계산이 되지 않도록 하기 위해서 매 반복마다 명시적으로 0으로 초기화가 필요함
• loss.backward()를 호출하여 prediction loss를 backpropagation 한다. 파이토치는 각 파라미터에 대한 loss의 gradient를 저장한다.
• optimizer.stop()를 호출하여 backward pass에서 계산된 gradient로 파라미터를 조정한다.

 

Full Implementation

위에서 살펴본 내용들을 종합하여, optimization loop를 다음과 같이 구현할 수 있습니다.

def train_loop(dataloader, model, loss_fn, optimizer):
    size = len(dataloader.dataset)
    for batch, (X, y) in enumerate(dataloader):
        # Compute prediction and loss
        pred = model(X)
        loss = loss_fn(pred, y)

        # Backpropagation
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if batch % 100 == 0:
            loss, current = loss.item(), batch * len(X)
            print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")


def test_loop(dataloader, model, loss_fn):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    test_loss, correct = 0, 0

    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")

이렇게 구현한 train_loop와 test_loop를 다음과 같이 사용하여, 학습 및 테스트를 수행합니다.

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!")

총 10번 반복하였고, 최종적으로 71.2%의 정확도를 달성했습니다.