MNIST 데이터셋 로드

학습 및 테스트 데이터를 준비하는 과정은 다음과 같다:

import torch
import torchvision
from torch.utils.data import DataLoader

transform = torchvision.transforms.ToTensor()

training_dataset = torchvision.datasets.MNIST(
    root='./data',
    train=True,
    transform=transform,
    download=True
)

testing_dataset = torchvision.datasets.MNIST(
    root='./data',
    train=False,
    transform=transform,
    download=True
)

train_dataloader = DataLoader(training_dataset, batch_size=64, shuffle=True)
test_dataloader = DataLoader(testing_dataset, batch_size=1000, shuffle=False)

데이터 시각화

학습 데이터의 특정 샘플을 확인하기 위해 다음 코드를 사용할 수 있다:

import matplotlib.pyplot as plt

index = 4000
image_tensor = training_dataset.data[index]
label_value = training_dataset.targets[index]

plt.imshow(image_tensor, cmap='gray')
plt.title(f"Label: {label_value}")
plt.show()

신경망 모델 정의

기본적인 완전 연결 계층(Fully Connected Layer) 기반 모델은 아래와 같다:

import torch.nn as nn

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

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

모델 학습

학습 및 평가 루프는 다음과 같이 구성된다:

import torch.optim as optim

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

network = SimpleNet().to(device)
loss_function = nn.CrossEntropyLoss()
optimizer_func = optim.Adam(network.parameters(), lr=0.001)

for epoch in range(1, 6):
    network.train()
    for inputs, targets in train_dataloader:
        inputs, targets = inputs.to(device), targets.to(device)
        
        optimizer_func.zero_grad()
        outputs = network(inputs)
        loss = loss_function(outputs, targets)
        loss.backward()
        optimizer_func.step()

    network.eval()
    correct_predictions = 0
    total_samples = 0
    
    with torch.no_grad():
        for inputs, targets in test_dataloader:
            inputs, targets = inputs.to(device), targets.to(device)
            predictions = network(inputs).argmax(dim=1)
            correct_predictions += (predictions == targets).sum().item()
            total_samples += targets.size(0)
    
    accuracy = correct_predictions / total_samples
    print(f"Epoch {epoch}, Test Accuracy: {accuracy:.4f}")

옵티마이저 비교

SGD 옵티마이저

sgd_optimizer = optim.SGD(
    network.parameters(),
    lr=0.1,
    momentum=0.9,
    weight_decay=5e-4
)

결과 예시:

• Epoch 1: 0.9572
• Epoch 2: 0.9683
• Epoch 3: 0.9692
• Epoch 4: 0.9735
• Epoch 5: 0.9705

Adam 옵티마이저

adam_optimizer = optim.Adam(network.parameters(), lr=1e-3)

결과 예시:

• Epoch 1: 0.9464
• Epoch 2: 0.9615
• Epoch 3: 0.9687
• Epoch 4: 0.9714
• Epoch 5: 0.9729

CNN 기반 모델로 개선

합성곱 신경망(Convolutional Neural Network)을 적용하여 성능을 향상시킬 수 있다:

class ConvNet(nn.Module):
    def __init__(self):
        super(ConvNet, self).__init__()
        self.feature_extractor = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=5),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
            
            nn.Conv2d(32, 64, kernel_size=5),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2)
        )
        
        self.classifier_head = nn.Sequential(
            nn.Flatten(),
            nn.Linear(64 * 4 * 4, 128),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),
            nn.Linear(128, 10)
        )

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

해당 모델을 사용했을 때의 결과:

• Epoch 1: 0.9828
• Epoch 2: 0.9890
• Epoch 3: 0.9906
• Epoch 4: 0.9911
• Epoch 5: 0.9914