간단한 이미지 분류 프로젝트 완성하기

프로젝트를 작은 블록으로 나누기

이번 장에서는 FashionMNIST 이미지 분류를 하나의 미니 프로젝트로 완성합니다. 지금까지 배운 내용을 순서대로 연결하면 PyTorch 프로젝트의 기본 흐름은 다음과 같습니다.

1. 데이터 준비
2. 모델 정의
3. 학습 함수 작성
4. 평가 함수 작성
5. 실행 코드에서 반복 학습
6. 검증 성능이 가장 좋은 모델 저장
7. 테스트 데이터로 최종 확인
8. 저장한 모델로 예측 확인

처음부터 파일을 여러 개로 나누기 어렵다면 train.py 하나로 시작해도 됩니다. 다만 코드를 쓸 때는 머릿속에서 위 블록을 구분해야 합니다. 그래야 오류가 났을 때 데이터 문제인지, 모델 문제인지, 학습 루프 문제인지 찾기 쉽습니다.

데이터 준비

FashionMNIST는 28×28 크기의 흑백 옷 이미지 데이터셋입니다. 입력 텐서 모양은 보통 [배치 크기, 1, 28, 28]입니다.

이번 프로젝트에서는 데이터를 세 부분으로 나누어 사용합니다.

• 학습 데이터: 모델 파라미터를 업데이트하는 데 사용합니다.

• 검증 데이터: 학습 중 어떤 모델이 가장 좋은지 고르는 데 사용합니다.

• 테스트 데이터: 모든 선택이 끝난 뒤 최종 성능을 확인하는 데 사용합니다.

테스트 데이터를 모델 선택에 반복해서 사용하면 실제 성능을 과대평가할 수 있습니다. 그래서 학습용 FashionMNIST 일부를 검증용으로 나누고, 원래 테스트 데이터는 마지막에만 확인합니다.

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

batch_size = 64

device = torch.device(
    "cuda" if torch.cuda.is_available()
    else "mps" if torch.backends.mps.is_available()
    else "cpu"
)

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

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

train_size = int(len(full_train_data) * 0.9)
val_size = len(full_train_data) - train_size

train_data, val_data = random_split(
    full_train_data,
    [train_size, val_size],
    generator=torch.Generator().manual_seed(42)
)

train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_data, batch_size=batch_size, shuffle=False)
test_loader = DataLoader(test_data, batch_size=batch_size, shuffle=False)

ToTensor()는 이미지를 텐서로 바꾸고, Normalize()는 값의 범위를 학습하기 좋은 형태로 맞춥니다. 여기서는 (0.5,)를 평균과 표준편차로 사용해 흑백 이미지 값을 대략 -1~1 범위로 바꿉니다.

완전연결 신경망과 CNN 비교하기

완전연결 신경망은 이미지를 한 줄로 펼쳐서 처리합니다. 구조가 단순하지만, 이미지의 위아래·좌우 관계를 직접 활용하지 못합니다.

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)

반면 CNN은 작은 필터를 이미지 위에서 움직이며 패턴을 찾습니다. 선, 모서리, 질감처럼 가까운 픽셀 사이의 관계를 보기 때문에 이미지 분류에 더 잘 맞습니다.

class SmallCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 16, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(16, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(32 * 7 * 7, 10)
        )

    def forward(self, x):
        x = self.features(x)
        return self.classifier(x)

처음에는 MLP()로 실행해 보고, 그다음 SmallCNN()으로 바꿔 정확도를 비교해 보세요. 이 장의 실행 코드는 CNN을 기준으로 이어 갑니다.

학습 함수와 평가 함수 만들기

학습 함수는 train() 모드에서 실행합니다. 평가 함수는 eval() 모드와 torch.no_grad()를 함께 사용합니다.

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

    for x, y in loader:
        x, y = x.to(device), y.to(device)

        pred = model(x)
        loss = loss_fn(pred, y)

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

        total_loss += loss.item() * x.size(0)
        correct += (pred.argmax(1) == y).sum().item()
        total += y.size(0)

    return total_loss / total, correct / total

@torch.no_grad()
def evaluate(model, loader, loss_fn, device):
    model.eval()
    total_loss, correct, total = 0, 0, 0

    for x, y in loader:
        x, y = x.to(device), y.to(device)
        pred = model(x)
        loss = loss_fn(pred, y)

        total_loss += loss.item() * x.size(0)
        correct += (pred.argmax(1) == y).sum().item()
        total += y.size(0)

    return total_loss / total, correct / total

손실을 더할 때 loss.item() * x.size(0)을 사용하는 이유는 마지막 배치 크기가 다를 수 있기 때문입니다. 샘플 수까지 반영해 누적한 뒤 전체 개수로 나누면 평균 손실을 더 정확하게 계산할 수 있습니다.

검증 성능이 가장 좋은 모델 저장하기

학습 중에는 검증 정확도가 가장 높을 때의 모델만 저장합니다. 테스트 정확도는 마지막에 한 번 확인하는 용도로 남겨 둡니다.

model = SmallCNN().to(device)
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

epochs = 10
best_val_acc = 0.0
history = {
    "train_loss": [], "val_loss": [],
    "train_acc": [], "val_acc": []
}

for epoch in range(epochs):
    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["train_loss"].append(train_loss)
    history["val_loss"].append(val_loss)
    history["train_acc"].append(train_acc)
    history["val_acc"].append(val_acc)

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

    print(
        f"epoch {epoch + 1}: "
        f"train loss={train_loss:.4f}, train acc={train_acc:.4f}, "
        f"val loss={val_loss:.4f}, val acc={val_acc:.4f}"
    )

학습 손실과 검증 손실이 함께 줄면 잘 배우고 있는 것입니다. 학습 정확도만 계속 오르고 검증 정확도가 멈추거나 떨어지면 과대학습일 수 있습니다. 둘 다 낮은 수준에 머무르면 모델이 너무 작거나 학습이 부족한 과소학습일 수 있습니다.

기록을 그래프로 확인하기

숫자만 보는 것보다 그래프를 보면 흐름이 더 잘 보입니다.

import matplotlib.pyplot as plt

plt.plot(history["train_loss"], label="train loss")
plt.plot(history["val_loss"], label="val loss")
plt.legend()
plt.show()

plt.plot(history["train_acc"], label="train acc")
plt.plot(history["val_acc"], label="val acc")
plt.legend()
plt.show()

검증 그래프까지 확인한 뒤, 저장해 둔 best model을 불러와 테스트 데이터에서 최종 성능을 측정합니다.

model = SmallCNN().to(device)
state_dict = torch.load(
    "best_fashion_cnn.pth",
    map_location=device,
    weights_only=True,
)
model.load_state_dict(state_dict)

test_loss, test_acc = evaluate(model, test_loader, loss_fn, device)
print(f"test loss={test_loss:.4f}, test acc={test_acc:.4f}")

저장한 모델로 예측 확인하기

마지막으로 저장한 모델을 사용해 테스트 이미지 몇 장을 예측해 봅니다.

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

images, labels = next(iter(test_loader))

model.eval()
with torch.no_grad():
    preds = model(images.to(device)).argmax(1).cpu()

for i in range(6):
    plt.subplot(2, 3, i + 1)

    # Normalize((0.5,), (0.5,))를 되돌려 화면에 보기 좋게 표시합니다.
    img = images[i].squeeze() * 0.5 + 0.5

    plt.imshow(img, cmap="gray")
    plt.title(f"pred: {classes[preds[i]]}ntrue: {classes[labels[i]]}")
    plt.axis("off")

plt.tight_layout()
plt.show()

정확도 숫자만 보면 어떤 종류의 이미지를 자주 틀리는지 알기 어렵습니다. 이렇게 샘플 예측을 눈으로 확인하면 모델이 헷갈리는 패턴을 더 쉽게 찾을 수 있습니다.

실험은 하나씩 바꾸기

모델 성능을 높이고 싶다면 한 번에 여러 값을 바꾸지 마세요. learning rate, batch size, epoch, 모델 크기 중 하나만 바꿔야 원인을 알 수 있습니다.

예를 들어 먼저 lr=0.001과 lr=0.0005를 비교합니다. 그다음 배치 크기 64와 128을 비교합니다. 이렇게 실험 기록을 남기면 PyTorch 코드는 단순한 예제가 아니라, 개선할 수 있는 프로젝트가 됩니다.