PyTorch 이미지 분류(Image Classification) 예제 (MNIST, CIFAR-10)

PyTorch 이미지 분류(Image Classification) 예제 (MNIST, CIFAR-10)

1. 서론

이미지 분류(Image Classification)는 딥러닝에서 가장 기본적인 문제 중 하나로, 입력된 이미지를 특정 클래스에 할당하는 작업입니다. 이번 글에서는 PyTorch를 사용하여 MNIST와 CIFAR-10 데이터셋을 활용한 이미지 분류 모델을 구현하는 방법을 소개합니다.

2. 데이터셋 소개

2.1 MNIST 데이터셋

MNIST는 손으로 쓴 숫자(0~9)로 구성된 28x28 픽셀의 흑백 이미지 데이터셋입니다. 총 60,000개의 학습 데이터와 10,000개의 테스트 데이터로 구성되어 있으며, 딥러닝 모델의 성능 평가에 자주 사용됩니다.

2.2 CIFAR-10 데이터셋

CIFAR-10은 10개의 일반적인 객체 클래스(비행기, 자동차, 새, 고양이, 사슴, 개, 개구리, 말, 배, 트럭)로 구성된 컬러 이미지 데이터셋입니다. 각 이미지는 32x32 픽셀 크기이며, 50,000개의 학습 데이터와 10,000개의 테스트 데이터로 구성되어 있습니다.

3. PyTorch를 활용한 이미지 분류 모델 구현

3.1 라이브러리 및 데이터 로딩

먼저 필요한 라이브러리를 불러오고, 데이터셋을 로드합니다.

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms

# 하이퍼파라미터 설정
batch_size = 64
learning_rate = 0.001
epochs = 10

# 데이터 전처리 변환 정의
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))  # MNIST의 경우 흑백 정규화 적용
])

# MNIST 데이터 로드
train_dataset = torchvision.datasets.MNIST(root='./data', train=True, transform=transform, download=True)
test_dataset = torchvision.datasets.MNIST(root='./data', train=False, transform=transform, download=True)

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

3.2 CNN 모델 정의

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(64 * 7 * 7, 128)
        self.fc2 = nn.Linear(128, 10)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.5)

    def forward(self, x):
        x = self.relu(self.conv1(x))
        x = self.pool(self.relu(self.conv2(x)))
        x = x.view(x.size(0), -1)
        x = self.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.fc2(x)
        return x

3.3 모델 학습

# 모델, 손실 함수, 옵티마이저 정의
model = CNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

def train(model, train_loader, criterion, optimizer, epochs):
    model.train()
    for epoch in range(epochs):
        running_loss = 0.0
        for images, labels in train_loader:
            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()
        print(f'Epoch [{epoch+1}/{epochs}], Loss: {running_loss/len(train_loader):.4f}')

train(model, train_loader, criterion, optimizer, epochs)

3.4 모델 평가

def evaluate(model, test_loader):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in test_loader:
            outputs = model(images)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    print(f'Accuracy: {100 * correct / total:.2f}%')

evaluate(model, test_loader)

4. CIFAR-10을 위한 모델 수정

CIFAR-10은 컬러 이미지이므로 in_channels=3으로 변경하고, Fully Connected Layer 크기를 조정해야 합니다.

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

train_dataset_cifar10 = torchvision.datasets.CIFAR10(root='./data', train=True, transform=transform_cifar10, download=True)
test_dataset_cifar10 = torchvision.datasets.CIFAR10(root='./data', train=False, transform=transform_cifar10, download=True)

train_loader_cifar10 = torch.utils.data.DataLoader(train_dataset_cifar10, batch_size=batch_size, shuffle=True)
test_loader_cifar10 = torch.utils.data.DataLoader(test_dataset_cifar10, batch_size=batch_size, shuffle=False)

class CIFAR10_CNN(nn.Module):
    def __init__(self):
        super(CIFAR10_CNN, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(64 * 8 * 8, 128)
        self.fc2 = nn.Linear(128, 10)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.5)

    def forward(self, x):
        x = self.relu(self.conv1(x))
        x = self.pool(self.relu(self.conv2(x)))
        x = x.view(x.size(0), -1)
        x = self.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.fc2(x)
        return x

5. 결론

이번 글에서는 MNIST와 CIFAR-10을 활용하여 이미지 분류 모델을 구현하는 방법을 살펴보았습니다. 기본적인 CNN 모델을 적용하는 방법을 익히고, 데이터셋의 특성에 맞게 모델을 수정하는 방법을 실습하였습니다.