PyTorch Loss 및 Optimizer 정리

PyTorch를 사용하여 모델 학습 시 중요한 두 가지 개념은 **손실 함수(Loss Function)**와 **최적화 함수(Optimizer)**입니다. 이 두 가지가 어떤 역할을 하고, 어떻게 사용하는지 간단히 정리해 보겠습니다.

---

1. 손실 함수 (Loss Function)

손실 함수는 모델의 출력과 실제 값 간의 차이를 계산합니다. 예를 들어, 아래는 평균 절대 오차(Mean Absolute Error, MAE)를 정의한 식과 이를 PyTorch에서 클래스로 구현한 코드입니다:

평균 절대 오차 (MAE)

PyTorch에서 MAE 구현

class Loss(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, input, target):
        return (1 / len(input)) * torch.sum(torch.abs(input - target))

 
여기에서 super().__init__()은 Loss 클래스의 부모클래스인 torch.nn.Module()의 __init__() 메서드를 호출하는 역할을 한다.
 
즉 super()는 Loss 클래스와의 관계성 내에서 생성된 부모 torch.nn.Module()일 것이다.
예시 :

class Parent:
    def __init__(self):
        print("Parent initialized")

class Child(Parent):
    def __init__(self):
        super().__init__()  # Calls Parent's __init__
        print("Child initialized")

c = Child()

Parent initialized
Child initialized

2. 손실 함수의 역할과 .backward()

• loss.backward() 호출 시:

input tensor의 grad 속성에 손실 함수의 **기울기(gradient)**가 저장됩니다.

주의해야 할 점은, criterion - 그리고 hence loss는 모델 파라미터가 어디있는지 모릅니다. 다만 input tensor의 위치와 그녀석의 gradient만 알고있죠.

 

물론 손실 계산을 loss = criterion(model(input), target)과 같은 식으로 하면, 입력텐서인 model(input)에 inherently 모델의 레이어 파라미터들이 개입될테니 이 경우 '간접적으로 알게되었다'고 볼 수 있겠죠. 그리고 gradient 또한 각각 모델 파라미터에 대응해서 저장될겁니다.

 

---

3. Optimizer(최적화 함수)

기울기를 통해 모델의 파라미터를 업데이트합니다. 대표적으로 **Stochastic Gradient Descent (SGD)**는 아래와 같이 모델의 파라미터를 업데이트합니다:

PyTorch에서 SGD 사용 예시

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

 

여기서 model.parameters()를 명시적으로 넘겨주는거에 주목! optimizer는 모델 파라미터(가중치, 편향)들이 어디 저장되어있는지 직접 알고있습니다. 

 

2에서 언급했듯 loss = criterion(model(input), output)과 같은 식으로 정의했다면, model(input)을 통해 loss는 inherently 모델 파라미터들 각각에 대응하는 손실

를 가지고있을겁니다. 

 

Optimizer가 하는 역할은, loss가 어떻게 그래디언트를 계산했건과는 상관없이 '그냥 모델 파라미터에 접착된 grad를 이용해 파라미터를 업데이트'해주는것 뿐입니다. 이미 model.parameters()에 grad가 접착되어있어야 하는거죠.

• optimizer.zero_grad(): 기울기를 초기화합니다. (PyTorch는 기본적으로 기울기를 누적합니다.)
• optimizer.step(): 계산된 기울기를 기반으로 파라미터를 업데이트합니다.

4. 학습 과정의 전체 흐름

학습 과정에서 손실 함수와 최적화 함수가 어떻게 작동하는지 아래 다이어그램으로 나타낼 수 있습니다:

학습 흐름

1. Input: 모델의 입력 데이터.
2. Target: 실제 값.
3. Criterion (손실 함수): 손실 계산.
4. Backward (역전파): loss.backward() 호출로 기울기 계산.
5. Optimizer:
   
   • optimizer.zero_grad(): 누적된 기울기 초기화.
   • optimizer.step(): 모델의 파라미터를 업데이트.

import torch
import torch.nn as nn
import torch.optim as optim
from tqdm import tqdm

# Loss and Optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

# Training Loop
def train_model():
    for epoch in tqdm(range(num_epochs)):
        model.train()
        train_loss, train_correct = 0, 0

        for images, labels in train_loader:
            images, labels = images.cuda(), labels.cuda()

            # Forward pass
            outputs = model(images)
            loss = criterion(outputs, labels)
            train_loss += loss.item()

            # Backward pass
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            # Accuracy
            _, preds = torch.max(outputs, 1)
            train_correct += (preds == labels).sum().item()

        val_loss, val_correct = 0, 0
        model.eval()
        with torch.no_grad():
            for images, labels in val_loader:
                images, labels = images.cuda(), labels.cuda()
                outputs = model(images)
                loss = criterion(outputs, labels)
                val_loss += loss.item()
                _, preds = torch.max(outputs, 1)
                val_correct += (preds == labels).sum().item()

        print(
            f"Epoch {epoch+1}/{num_epochs}, "
            f"Train Loss: {train_loss/len(train_loader):.4f}, Train Acc: {train_correct/len(train_dataset):.4f}, "
            f"Val Loss: {val_loss/len(val_loader):.4f}, Val Acc: {val_correct/len(val_dataset):.4f}"
        )

train_model()