PyTorch를 사용한 딥러닝 모델 최적화: 훈련 루프 구현

모델과 데이터셋을 준비한 후에는 모델의 파라미터를 최적화하여 학습시키는 과정이 필요합니다. 이 과정은 반복적인 훈련 절차로 구성되며, 각 반복 단계에서 모델은 입력 데이터에 대한 예측을 수행하고, 실제 정답과의 오차를 계산한 후 그 기울기를 기반으로 가중치를 갱신합니다. 이러한 학습 프로세스는 전방 전파(forward pass), 손실 계산, 역전파(backward pass), 그리고 옵티마이저를 통한 파라미터 업데이트로 이루어집니다.

기초 설정

먼저 FashionMNIST 데이터셋을 불러오고, 이를 미니배치로 나누어 처리할 수 있도록 DataLoader를 구성합니다. 또한 간단한 피드포워드 신경망을 정의합니다.

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

# 데이터 전처리 및 로드
transform = transforms.Compose([transforms.ToTensor()])

train_dataset = datasets.FashionMNIST(
    root='data', train=True, download=True, transform=transform
)
test_dataset = datasets.FashionMNIST(
    root='data', train=False, download=True, transform=transform
)

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

# 신경망 정의
class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.flatten = nn.Flatten()
        self.feature_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)
        output = self.feature_stack(x)
        return output

model = SimpleNet()

하이퍼파라미터 설정

학습 과정을 제어하기 위한 하이퍼파라미터를 지정합니다.

  • 에포크(Epoch): 전체 훈련 데이터를 한 번 순회하는 횟수
  • 배치 크기(Batch Size): 한 번의 기울기 갱신에 사용되는 샘플 수
  • 학습률(Learning Rate): 파라미터 갱신의 크기 조절 값
lr = 0.001
batch_size = 64
n_epochs = 10

손실 함수와 옵티마이저

분류 문제에서는 교차 엔트로피 손실(CrossEntropyLoss)을 주로 사용합니다. 이 손실 함수는 소프트맥스 활성화 함수와 네거티브 로그 가능도(NLL)를 결합한 형태입니다. 옵티마이저로는 확률적 경사하강법(SGD)을 사용합니다.

# 손실 함수
criterion = nn.CrossEntropyLoss()

# 옵티마이저
optimizer = torch.optim.SGD(model.parameters(), lr=lr)

훈련 및 평가 루프

모델의 학습과 성능 평가를 위한 두 개의 함수를 정의합니다. 훈련 루프에서는 기울기를 초기화하고 역전파를 수행한 후 옵티마이저로 가중치를 갱신합니다. 평가 루프에서는 기울기 계산을 비활성화하여 메모리 사용을 줄입니다.

def train_epoch(dataloader, model, loss_fn, opt):
    size = len(dataloader.dataset)
    model.train()  # 훈련 모드 설정
    for batch_idx, (X_batch, y_batch) in enumerate(dataloader):
        # 예측 및 손실 계산
        pred = model(X_batch)
        loss = loss_fn(pred, y_batch)

        # 역전파
        opt.zero_grad()
        loss.backward()
        opt.step()

        if batch_idx % 100 == 0:
            current_loss = loss.item()
            current_step = batch_idx * len(X_batch)
            print(f"Train Loss: {current_loss:.6f} [{current_step:>5d}/{size:>5d}]")

def evaluate_model(dataloader, model, loss_fn):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    model.eval()  # 평가 모드 설정
    total_loss = 0
    correct_count = 0

    with torch.no_grad():
        for X_batch, y_batch in dataloader:
            pred = model(X_batch)
            total_loss += loss_fn(pred, y_batch).item()
            correct_count += (pred.argmax(dim=1) == y_batch).type(torch.float).sum().item()

    avg_loss = total_loss / num_batches
    accuracy = correct_count / size
    print(f"Evaluation Results:\nAccuracy: {(accuracy*100):>0.1f}%, Avg Loss: {avg_loss:>8f}\n")

전체 학습 실행

지정된 에포크 동안 훈련과 평가를 반복합니다.

for epoch in range(n_epochs):
    print(f"Epoch [{epoch+1}/{n_epochs}]")
    print("-" * 30)
    train_epoch(train_loader, model, criterion, optimizer)
    evaluate_model(test_loader, model, criterion)

print("Training completed.")

출력 결과는 각 에포크마다 훈련 손실과 테스트 정확도를 보여주며, 시간이 지남에 따라 모델의 성능이 점차 개선됨을 확인할 수 있습니다.

태그: PyTorch 딥러닝 신경망 학습 손실함수 옵티마이저

7월 4일 19:51에 게시됨