Skip to main content
Summary

이 장에서 배울 것

앞 장에서는 학습률, 옵티마이저, 정규화, train()과 eval() 모드의 의미를 배웠다. 이번 장에서는 그 개념을 PyTorch 코드로 연결한다. 목표는 단순하다. 데이터를 준비하고, 모델을 만들고, 손실을 계산하고, 역전파를 실행하고, 가중치를 업데이트하고, 평가까지 수행하는 전체 학습 루프를 직접 작성하는 것이다.

이 장을 마치면 다음 흐름을 코드로 설명할 수 있어야 한다.

optimizer.zero_grad()
outputs = model(inputs)
loss = loss_fn(outputs, labels)
loss.backward()
optimizer.step()

이 다섯 줄은 PyTorch 학습의 핵심이다. 짧지만 순서가 중요하다. 한 줄만 빠지거나 위치가 바뀌어도 학습이 느려지거나, 아예 되지 않거나, 이상한 결과가 나온다.

PyTorch 학습 루프

왜 이 주제가 중요한가

Keras 같은 고수준 도구에서는 fit() 한 번으로 학습이 실행된다. 반면 PyTorch는 학습 루프를 직접 쓰는 방식이 기본이다. 번거롭게 보이지만 장점이 있다. 데이터가 어디에서 들어오는지, 손실이 언제 계산되는지, 기울기가 언제 초기화되고 업데이트되는지 눈으로 확인할 수 있다.

PyTorch 공식 튜토리얼도 Dataset, DataLoader, 손실 함수, 옵티마이저를 조합해 하나의 학습 루프로 만드는 흐름을 강조한다.1 즉 PyTorch를 제대로 쓴다는 것은 모델 구조만 아는 것이 아니라 학습 루프의 상태 변화를 이해한다는 뜻이다.

학습 루프는 딥러닝의 심장박동과 같다. 한 번의 박동은 하나의 미니배치에서 일어난다. 데이터를 넣고, 예측하고, 틀린 정도를 재고, 그 틀림을 줄이는 방향으로 가중치를 조금 바꾼다. 이 과정을 수천 번 반복하면서 모델이 학습된다.

핵심 개념 설명

Tensor: PyTorch가 계산하는 숫자 묶음

Tensor는 PyTorch에서 데이터를 담는 기본 단위다. 숫자 하나일 수도 있고, 벡터, 행렬, 이미지 묶음일 수도 있다. 예를 들어 Fashion-MNIST 이미지 한 장은 보통 [1, 28, 28] 모양의 Tensor다. 여기서 1은 채널 수, 28, 28은 높이와 너비다.

학습 루프에서 Tensor를 볼 때는 세 가지를 확인해야 한다.

  • shape: 모양이 모델 입력과 맞는가?

  • dtype: 자료형이 손실 함수와 맞는가?

  • device: CPU와 GPU 중 같은 장치에 있는가?

초보자의 많은 오류가 이 세 가지 중 하나에서 나온다.

Dataset과 DataLoader: 데이터를 배치로 꺼내는 장치

Dataset은 데이터 하나를 어떻게 꺼낼지 정의한다. DataLoader는 Dataset에서 여러 개를 모아 미니배치로 만든다. PyTorch 문서에서는 Dataset이 개별 샘플 접근을 담당하고, DataLoader가 샘플을 배치로 묶어 학습 루프에 제공한다고 설명한다.1

train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

이 코드는 학습 데이터를 64개씩 묶어서 꺼내겠다는 뜻이다. shuffle=True는 매 에폭마다 데이터 순서를 섞는다. 학습 데이터는 섞는 것이 일반적이다. 평가 데이터는 보통 섞지 않는다.

nn.Module: 모델을 담는 기본 틀

PyTorch 모델은 보통 nn.Module을 상속해서 만든다. __init__()에는 사용할 층을 정의하고, forward()에는 입력이 출력으로 변하는 흐름을 적는다.

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Flatten(),
            nn.Linear(28 * 28, 128),
            nn.ReLU(),
            nn.Linear(128, 10)
        )

    def forward(self, x):
        return self.net(x)

여기서 출력은 10개 숫자다. Fashion-MNIST에는 10개 클래스가 있기 때문이다. 이 숫자를 로짓(logit)이라고 부른다. 로짓은 아직 확률로 바뀌기 전의 점수다.

손실 함수: 예측이 얼마나 틀렸는지 재는 함수

분류 문제에서는 nn.CrossEntropyLoss()를 자주 쓴다. 이 함수는 모델 출력 로짓과 정답 레이블을 비교해 손실을 계산한다. 주의할 점은 정답 레이블이 원-핫 벡터가 아니라 정수 클래스 번호여야 한다는 것이다.

loss_fn = nn.CrossEntropyLoss()

예를 들어 정답이 운동화라면 [0,0,0,0,0,0,0,1,0,0]이 아니라 7처럼 들어간다.

Optimizer: 기울기로 가중치를 바꾸는 객체

옵티마이저는 모델 파라미터를 받아서, 계산된 기울기를 바탕으로 값을 업데이트한다. PyTorch의 torch.optim은 SGD, Adam, AdamW 같은 여러 최적화 알고리즘을 제공한다.2

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

여기서 model.parameters()를 넘기는 이유는 명확하다. 옵티마이저에게 "이 파라미터들을 업데이트하라"고 알려주는 것이다.

train 모드와 eval 모드

model.train()은 모델을 학습 모드로 둔다. Dropout, BatchNorm처럼 학습과 평가 때 다르게 동작하는 층이 있다면 이 설정이 중요하다. model.eval()은 평가 모드다. 평가할 때는 보통 torch.no_grad()도 함께 사용한다. 평가에서는 가중치를 바꾸지 않으므로 기울기를 저장할 필요가 없다.3

model.train()      # 학습할 때
model.eval()       # 평가할 때

with torch.no_grad():
    outputs = model(inputs)

device: CPU와 GPU 위치 맞추기

GPU가 있으면 학습을 더 빠르게 할 수 있다. 다만 모델과 데이터가 같은 장치에 있어야 한다.

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

images = images.to(device)
labels = labels.to(device)

이 장에서는 장치 이동 패턴만 익힌다. 여러 GPU, mixed precision, 세부 성능 최적화는 이 책의 범위를 넘기므로 다음 권이나 부록에서 다룬다.

쉬운 비유 또는 사례

학습 루프를 시험 공부에 비유해 보자.

  • DataLoader: 문제집에서 오늘 풀 문제 묶음을 꺼낸다.

  • model(inputs): 문제를 풀고 답을 낸다.

  • loss_fn(outputs, labels): 채점해서 얼마나 틀렸는지 점수화한다.

  • loss.backward(): 어떤 개념을 얼마나 잘못 이해했는지 분석한다.

  • optimizer.step(): 다음에는 덜 틀리도록 공부 방향을 수정한다.

  • optimizer.zero_grad(): 이전 문제의 오답 분석을 지우고 새 문제를 본다.

중요한 점은 "오답 분석을 지우는 일"이다. PyTorch는 기본적으로 기울기를 누적한다. 그래서 매 배치마다 zero_grad()를 호출하지 않으면 이전 배치의 기울기까지 섞인다. 특별히 의도한 것이 아니라면 잘못된 학습이 된다.

직접 해보는 실습: Fashion-MNIST MLP 분류

이번 실습에서는 Fashion-MNIST 이미지를 MLP로 분류한다. MLP는 이미지를 1차원으로 펼쳐서 분류하는 가장 단순한 신경망이다. 이미지를 다루는 CNN에서는 CNN을 사용해 이미지의 공간 구조를 더 잘 활용하는 방법을 배운다.

이번 실습에서는 데이터를 세 부분으로 나눈다.

  • train set: 가중치를 학습하는 데이터

  • validation set: 모델 선택과 설정 비교에 사용하는 데이터

  • test set: 모든 선택이 끝난 뒤 마지막에 한 번 확인하는 데이터

이 구분은 중요하다. test set을 매 epoch 확인하며 가장 좋은 모델을 고르면, test set이 사실상 모델 선택에 사용된 것이 된다. 그러면 최종 성능을 공정하게 말하기 어려워진다.

1단계: 라이브러리와 환경 준비

import random
import numpy as np
import pandas as pd
import torch
from torch import nn
from torch.utils.data import DataLoader, random_split
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay

seed = 42
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)

if torch.cuda.is_available():
    torch.cuda.manual_seed_all(seed)

generator = torch.Generator().manual_seed(seed)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("device:", device)

실행 환경마다 결과가 조금 다를 수 있다. 하지만 같은 환경에서 seed를 고정하면 비교가 쉬워진다.

2단계: 데이터셋과 DataLoader 만들기

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

full_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_size = int(len(full_train_dataset) * 0.8)
val_size = len(full_train_dataset) - train_size
train_dataset, val_dataset = random_split(
    full_train_dataset,
    [train_size, val_size],
    generator=generator
)

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

classes = [
    "T-shirt/top", "Trouser", "Pullover", "Dress", "Coat",
    "Sandal", "Shirt", "Sneaker", "Bag", "Ankle boot"
]

images, labels = next(iter(train_loader))
print(images.shape)  # torch.Size([64, 1, 28, 28])
print(labels.shape)  # torch.Size([64])

images의 shape는 [배치 크기, 채널, 높이, 너비]다. MLP는 2차원 이미지를 직접 처리하지 못하므로 모델 안에서 nn.Flatten()으로 펼친다.

3단계: 모델 정의

class FashionMLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Flatten(),
            nn.Linear(28 * 28, 256),
            nn.ReLU(),
            nn.Dropout(p=0.2),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, 10)
        )

    def forward(self, x):
        return self.net(x)

model = FashionMLP().to(device)
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)

Dropout과 weight_decay는 앞 장에서 배운 정규화 기법이다. 여기서는 과적합을 조금 줄이기 위해 사용한다.

4단계: 학습 함수 작성

def train_one_epoch(model, dataloader, loss_fn, optimizer, device):
    model.train()
    total_loss = 0.0
    correct = 0
    total = 0

    for images, labels in dataloader:
        images = images.to(device)
        labels = labels.to(device)

        # 1. 이전 배치의 기울기 초기화
        optimizer.zero_grad()

        # 2. 순전파: 예측값 계산
        outputs = model(images)

        # 3. 손실 계산
        loss = loss_fn(outputs, labels)

        # 4. 역전파: 기울기 계산
        loss.backward()

        # 5. 파라미터 업데이트
        optimizer.step()

        batch_size = images.size(0)
        total_loss += loss.item() * batch_size
        preds = outputs.argmax(dim=1)
        correct += (preds == labels).sum().item()
        total += batch_size

    avg_loss = total_loss / total
    accuracy = correct / total
    return avg_loss, accuracy

loss.item()은 Tensor에 담긴 숫자를 파이썬 숫자로 꺼낸다. 로그를 저장할 때는 보통 loss 자체가 아니라 loss.item()을 사용한다.

5단계: 평가 함수 작성

def evaluate(model, dataloader, loss_fn, device):
    model.eval()
    total_loss = 0.0
    correct = 0
    total = 0
    all_preds = []
    all_labels = []

    with torch.no_grad():
        for images, labels in dataloader:
            images = images.to(device)
            labels = labels.to(device)

            outputs = model(images)
            loss = loss_fn(outputs, labels)

            batch_size = images.size(0)
            total_loss += loss.item() * batch_size
            preds = outputs.argmax(dim=1)
            correct += (preds == labels).sum().item()
            total += batch_size

            all_preds.extend(preds.cpu().numpy())
            all_labels.extend(labels.cpu().numpy())

    avg_loss = total_loss / total
    accuracy = correct / total
    return avg_loss, accuracy, all_preds, all_labels

평가 함수에는 backward()와 optimizer.step()이 없다. 평가의 목적은 가중치를 바꾸는 것이 아니라 현재 모델이 얼마나 잘하는지 확인하는 것이다.

6단계: 전체 학습 실행과 로그 저장

epochs = 5
history = []
best_val_acc = 0.0

for epoch in range(1, epochs + 1):
    train_loss, train_acc = train_one_epoch(
        model, train_loader, loss_fn, optimizer, device
    )
    val_loss, val_acc, _, _ = evaluate(
        model, val_loader, loss_fn, device
    )

    history.append({
        "epoch": epoch,
        "train_loss": train_loss,
        "train_acc": train_acc,
        "val_loss": val_loss,
        "val_acc": val_acc
    })

    print(
        f"Epoch {epoch:02d} | "
        f"train_loss={train_loss:.4f}, train_acc={train_acc:.4f} | "
        f"val_loss={val_loss:.4f}, val_acc={val_acc:.4f}"
    )

    if val_acc > best_val_acc:
        best_val_acc = val_acc
        torch.save(model.state_dict(), "best_fashion_mlp.pth")

pd.DataFrame(history).to_csv("training_log.csv", index=False)
print("best_val_acc:", best_val_acc)

출력 예시는 다음과 비슷하다. 정확한 숫자는 환경과 버전에 따라 달라질 수 있다.

Epoch 01 | train_loss=0.6123, train_acc=0.7820 | val_loss=0.4568, val_acc=0.8351
Epoch 02 | train_loss=0.4315, train_acc=0.8459 | val_loss=0.4162, val_acc=0.8503
Epoch 03 | train_loss=0.3894, train_acc=0.8599 | val_loss=0.3975, val_acc=0.8588

7단계: test set으로 최종 평가하고 혼동 행렬 저장

모델 선택은 validation set으로 끝냈다. 이제 저장해 둔 가장 좋은 모델을 불러와 test set에서 마지막으로 평가한다.

model.load_state_dict(torch.load("best_fashion_mlp.pth", map_location=device))
test_loss, test_acc, preds, labels = evaluate(model, test_loader, loss_fn, device)
print(f"test_loss={test_loss:.4f}, test_acc={test_acc:.4f}")

cm = confusion_matrix(labels, preds)
disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=classes)

fig, ax = plt.subplots(figsize=(10, 10))
disp.plot(ax=ax, xticks_rotation=45, cmap="Blues")
plt.title(f"Fashion-MNIST Confusion Matrix | test_acc={test_acc:.4f}")
plt.tight_layout()
plt.savefig("confusion_matrix.png", dpi=150)
plt.show()

실습 결과물은 세 가지다.

  • training_log.csv: 에폭별 train/validation 손실과 정확도

  • best_fashion_mlp.pth: validation 성능이 가장 좋았던 모델 가중치

  • confusion_matrix.png: test set에서의 오류 유형을 보여주는 그림

이 세 파일을 보면 단순히 "학습했다"가 아니라 "어떻게 학습했고, 어디서 틀렸는지"까지 설명할 수 있다.

흔한 실수와 오류 분석

1. shape mismatch

대표 오류는 다음과 같다.

RuntimeError: mat1 and mat2 shapes cannot be multiplied

대부분 Linear 층에 들어가는 입력 크기가 맞지 않을 때 발생한다. Fashion-MNIST 이미지는 [batch, 1, 28, 28]이다. 이를 nn.Linear(28*28, 256)에 넣으려면 먼저 [batch, 784]로 펼쳐야 한다.

해결 방법은 nn.Flatten()을 모델 앞부분에 넣는 것이다.

nn.Flatten()

또는 forward() 안에서 다음처럼 처리할 수도 있다.

x = x.view(x.size(0), -1)

2. dtype 오류

대표 오류는 다음과 같다.

RuntimeError: expected scalar type Long but found Float

nn.CrossEntropyLoss()의 정답 레이블은 정수형 클래스 번호여야 한다. 즉 labels는 보통 torch.long 타입이어야 한다. 직접 만든 데이터셋에서 이런 오류가 난다면 다음처럼 확인한다.

print(labels.dtype)
labels = labels.long()

단, Fashion-MNIST처럼 torchvision이 제공하는 분류 데이터셋은 보통 레이블 타입이 이미 맞게 들어온다.

3. device 오류

대표 오류는 다음과 같다.

Expected all tensors to be on the same device

모델은 GPU에 있는데 데이터는 CPU에 있거나, 반대일 때 발생한다.

model = model.to(device)
images = images.to(device)
labels = labels.to(device)

모델과 입력, 정답을 모두 같은 device로 옮겨야 한다.

4. loss가 줄지 않는 문제

손실이 거의 줄지 않는다면 다음을 확인한다.

  • optimizer.zero_grad()를 호출했는가?

  • loss.backward() 뒤에 optimizer.step()을 호출했는가?

  • 학습률이 너무 크거나 너무 작지 않은가?

  • 입력 정규화를 했는가?

  • 모델 출력에 softmax를 붙인 뒤 CrossEntropyLoss에 넣고 있지 않은가?

nn.CrossEntropyLoss()는 내부적으로 로짓을 처리한다. 따라서 일반적인 다중 분류에서는 마지막 층 뒤에 Softmax를 직접 붙이지 않는다.

5. 평가 결과가 불안정한 문제

평가할 때 model.eval()을 빼먹으면 Dropout 같은 층이 계속 학습 모드로 동작한다. 그러면 같은 데이터에 대해서도 결과가 흔들릴 수 있다.

model.eval()
with torch.no_grad():
    ...

이 두 줄은 평가 루프의 기본 습관으로 만들자.

6. test set으로 모델을 고르는 문제

test set은 최종 확인용이다. 매 epoch test accuracy를 보고 가장 좋은 모델을 저장하면, 이미 test set을 모델 선택에 사용한 것이다. 실험 중 모델을 고를 때는 validation set을 사용하고, test set은 모든 선택이 끝난 뒤 마지막에 한 번만 확인하는 습관을 들이자.

요약

PyTorch 학습 루프는 다음 순서로 움직인다.

  1. DataLoader가 미니배치를 제공한다.

  2. 입력과 정답을 device로 옮긴다.

  3. optimizer.zero_grad()로 이전 기울기를 초기화한다.

  4. model(inputs)로 예측한다.

  5. loss_fn(outputs, labels)로 손실을 계산한다.

  6. loss.backward()로 기울기를 계산한다.

  7. optimizer.step()으로 가중치를 업데이트한다.

  8. 평가할 때는 model.eval()과 torch.no_grad()를 사용한다.

  9. 모델 선택은 validation set으로 하고, test set은 마지막 평가에 사용한다.

이 루프를 이해하면 앞에서 배운 신경망, 손실 함수, 역전파, 최적화가 하나의 코드 흐름으로 연결된다.

퀴즈

  1. optimizer.zero_grad()를 매 배치마다 호출하는 이유는 무엇인가?

  2. loss.backward()와 optimizer.step()은 각각 어떤 일을 하는가?

  3. model.train()과 model.eval()의 차이는 무엇인가?

  4. nn.CrossEntropyLoss()에 들어가는 모델 출력과 정답 레이블의 모양은 보통 어떻게 되는가?

  5. GPU를 사용할 때 모델과 Tensor에 공통으로 해주어야 하는 작업은 무엇인가?

  6. validation set과 test set을 구분해야 하는 이유는 무엇인가?

정답 예시:

  1. PyTorch는 기울기를 기본적으로 누적하므로 이전 배치의 기울기를 지우기 위해서다.

  2. backward()는 기울기를 계산하고, step()은 그 기울기로 파라미터를 업데이트한다.

  3. 학습 모드와 평가 모드에서 Dropout, BatchNorm 같은 층의 동작이 달라진다.

  4. 출력은 [batch_size, num_classes] 로짓, 정답은 [batch_size] 정수 레이블이다.

  5. 둘 다 같은 device로 이동해야 한다.

  6. validation set은 모델 선택용이고, test set은 최종 일반화 성능을 확인하는 용도이기 때문이다.

체크리스트

  • Dataset과 DataLoader의 역할을 구분할 수 있다.

  • nn.Module로 간단한 MLP를 정의할 수 있다.

  • zero_grad → forward → loss → backward → step 순서를 설명할 수 있다.

  • model.train()과 model.eval()을 올바른 위치에 넣을 수 있다.

  • CPU/GPU device 오류를 해결할 수 있다.

  • validation set으로 가장 좋은 모델을 저장할 수 있다.

  • test set은 마지막 평가에 사용해야 함을 이해했다.

  • 학습 로그를 CSV로 저장할 수 있다.

  • 혼동 행렬로 모델의 오류 유형을 확인할 수 있다.

다음 장 연결

이번 장의 MLP는 이미지를 한 줄로 펼쳐서 분류했다. 이 방식은 간단하지만 이미지의 중요한 구조를 잃는다. 예를 들어 신발의 밑창, 셔츠의 소매, 가방의 손잡이는 주변 픽셀과의 관계 속에서 의미가 생긴다. MLP는 이런 공간적 관계를 적극적으로 활용하지 못한다.

다음 장에서는 CNN을 배운다. CNN은 이미지를 펼치지 않고 작은 필터로 훑으면서 지역 패턴을 찾는다. 이번 장에서 만든 학습 루프는 그대로 사용할 수 있다. 달라지는 것은 주로 모델 구조다. 즉 앞으로도 학습의 기본 흐름은 같다. 데이터가 들어오고, 모델이 예측하고, 손실을 계산하고, 역전파로 가중치를 업데이트한다.

참고

1Training with PyTorch

2torch.optim 문서

3PyTorch 최적화 튜토리얼

책 목록

AI 교육 3권 : 딥러닝 핵심: 입력에서 업데이트까지, 신경망을 제대로 이해하는 법

전체 목차 보기

이 노트는 보도·비평·교육·연구 목적으로 작성되었습니다. 저작권 문의가 있으시면 에서 알려주세요.

틸노트는 사용자가 AI를 활용해 노트를 쉽게 작성하거나 편집할 수 있는 서비스입니다.

AI 스퀘어 단톡방에 함께해요

AI 활용과 틸노트 사용 팁을 나누는 오픈채팅방입니다. 입장 비밀번호는 tilnote입니다.

참여하기