[Python] PyTorch 활용해서 손글씨 데이터를 숫자로 분류하기

안녕하세요

오늘은 PyTorch 활용해서 MNIST 손글씨 데이터를, MLP로 분류하는 간단한 예제를 해보겠습니다.

 

우선 터미널에서

 

pip install torch torchvision

pip install certificate

 

을 설치해줍니다.

 

전체 코드입니다.

 

mnist_mlp.py

"""
MNIST 손글씨 숫자(0~9) 분류를 위한 '첫 PyTorch 예제'
- Dataset/DataLoader 쓰는 법
- nn.Module로 모델 만들기
- 손실함수/옵티마이저/학습 루프/평가까지 한 번에
"""
from typing import Self

import torch

# macOS 의 경우 torchvision의 다운로드가 사용하는 기본 HTTPS 컨텍스트에 certifi 인증서 번들을 물려주기위해 활성화할것
# import ssl, certifi
# ssl._create_default_https_context = lambda: ssl.create_default_context(cafile=certifi.where())

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


# 1) 디바이스 자동 선택: CUDA(GPU)가 있으면 쓰고, 없으면 CPU 사용
device = "cuda" if torch.cuda.is_available() else "cpu"
print("device:", device)

# 2) 데이터 전처리 파이프라인: 이미지를 Tensor로 변환(0~1 스케일)
#   * 여기서는 간단하게 ToTensor만 씀 (정규화는 생략해도 기본 성능 잘 나옴)
transform = transforms.ToTensor()

# 3) 학습/테스트 데이터셋 준비 (최초 1회 자동 다운로드)
train_ds = datasets.MNIST(root="./data", train=True, download=True, transform=transform)
test_ds  = datasets.MNIST(root="./data", train=False, download=True, transform=transform)

# 4) 배치 단위로 데이터를 꺼내주는 DataLoader
#   - shuffle=True: 매 epoch마다 데이터 순서를 섞어서 일반화에 도움
train_loader = DataLoader(train_ds, batch_size=64, shuffle=True)
test_loader  = DataLoader(test_ds, batch_size=64, shuffle=False)


# 5) 모델 정의: 아주 단순한 MLP (이미지를 1D로 펴서 분류)
class SimpleMLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Flatten(),             # (B, 1, 28, 28) -> (B, 784)
            nn.Linear(28*28, 256),    # 784 -> 256
            nn.ReLU(),                # 비선형 활성화
            nn.Linear(256, 128),      # 256 -> 128
            nn.ReLU(),
            nn.Linear(128, 10)        # 128 -> 10 (클래스 로짓)
        )
    def forward(self, x):
        return self.net(x)


model = SimpleMLP().to(device)

# 6) 손실함수 & 옵티마이저
#   - CrossEntropyLoss: 다중 클래스 분류에서 표준적인 선택
#   - Adam: 학습이 안정적이고 세팅이 간단
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)

# 7) 학습 루프 (3 epoch면 충분히 90%+ 정확도 나옴)
for epoch in range(3):
    model.train()                     # 드롭아웃/배치정규화 등 '학습 모드'로 전환
    running_loss = 0.0
    running_correct = 0
    total = 0

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

        # (1) 순전파
        logits = model(images)        # 크기: (B, 10)
        loss = criterion(logits, labels)

        # (2) 역전파 & 가중치 갱신
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # 로그용 누적
        running_loss += loss.item() * images.size(0)
        preds = logits.argmax(dim=1)               # 예측 클래스
        running_correct += (preds == labels).sum().item()
        total += labels.size(0)

    train_loss = running_loss / total
    train_acc = 100.0 * running_correct / total
    print(f"Epoch {epoch+1} | loss: {train_loss:.4f} | acc: {train_acc:.2f}%")


# 8) 평가 (테스트 세트)
model.eval()                          # '평가 모드' 전환 (드롭아웃 등 비활성화)
test_correct, test_total = 0, 0
with torch.no_grad():                 # 평가 시엔 그래디언트 계산 불필요 -> 메모리/속도 절약
    for images, labels in test_loader:
        images, labels = images.to(device), labels.to(device)
        logits = model(images)
        preds = logits.argmax(dim=1)
        test_correct += (preds == labels).sum().item()
        test_total += labels.size(0)

print(f"Test Accuracy: {100.0 * test_correct / test_total:.2f}%")

# 9) 샘플 몇 개 눈으로 확인 (정답 vs 예측)
images, labels = next(iter(test_loader))
images, labels = images.to(device), labels.to(device)
with torch.no_grad():
    logits = model(images[:8])
    preds = logits.argmax(dim=1)

print("Ground truth:", labels[:8].tolist())
print("Predicted   :", preds[:8].tolist())

 

이렇게

 

28×28 픽셀짜리 손글씨 이미지 → [0~1] 범위의 숫자 데이터로 변환 → MLP(다층 퍼셉트론)가 0~9 중 하나로 분류
즉 학습 데이터를 통해 가중치를 조정해서, 새 이미지도 맞출 수 있게 하는 총 과정입니다.

 

시연결과입니다.