본문 바로가기
카테고리 없음

PyTorch로 이미지 분류 모델 만들기 - 딥러닝 시작하기

by 북더기 2025. 4. 1.

딥러닝은 머신러닝의 하위 분야에 속하며, 인공신경망을 활용하여 사람의 뇌처럼 데이터를 학습하고 예측하는 모델을 만드는 기술입니다. 최근에는 이미지, 음성, 텍스트 등 다양한 분야에서 딥러닝이 활용되고 있으며, 그 중에서도 이미지 분류는 딥러닝 모델이 특히 강점을 보이는 분야 중 하나입니다. 이번 글에서는 PyTorch를 사용하여 간단한 이미지 분류 모델을 만드는 과정을 다루어보겠습니다. 실습 위주로 진행하며, 데이터셋 로딩부터 모델 학습, 평가까지 전체 흐름을 단계별로 살펴가 보겠습니다.

1. 데이터셋 준비와 전처리

딥러닝 모델을 학습시키기 위해 가장 먼저 해야 할 일은, 적절한 데이터셋을 준비하는 것입니다. PyTorch에서는 torchvision 라이브러리를 통해 다양한 이미지 데이터셋을 쉽게 불러올 수 있는 장점이 있습니다. 이 중에서 대표적으로는 CIFAR-10이나 MNIST 데이터셋이 많이 사용되며, 이번 예제에서는 간단한 이미지 분류 실습을 위해 CIFAR-10을 사용하여 코드로 구현해 보겠습니다.

아래의 예제로 함께 살펴보겠습니다.

import torch
import torchvision
import torchvision.transforms as transforms

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

trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=32, shuffle=True)

testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=32, shuffle=False)

위 코드에서는 CIFAR-10 데이터셋을 로딩하고, 정규화를 통해 입력 이미지의 픽셀 값을 0을 중심으로 스케일링합니다. 또한 DataLoader를 통해 데이터를 배치 단위로 불러올 수 있도록 설정하였습니다. 이러한 전처리 과정은 모델의 학습 성능에 큰 영향을 주기 때문에 신중히 설계하는 것이 중요한 부분입니다.

2. 신경망 모델 정의 및 학습

데이터를 준비한 다음에는 신경망 모델을 정의해야 합니다. 이 PyTorch에서는 nn.Module 클래스를 상속받아 모델 구조를 쉽게 구성할 수 있습니다. 여기서는 간단한 합성곱 신경망(CNN)을 구성하여 이미지를 분류하는 모델을 한번 만들어보겠습니다.

아래의 예제로 함께 살펴보겠습니다.

import torch.nn as nn
import torch.nn.functional as F

class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, 3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(16, 32, 3, padding=1)
        self.fc1 = nn.Linear(32 * 8 * 8, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 32 * 8 * 8)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

이 모델은, 두 개의 합성곱 층과 풀링층을 거쳐 완전 연결층으로 이어지는 구조를 가지고 있습니다. 이후에 모델을 GPU 또는 CPU에 올려 학습을 진행합니다. 학습 시에는 손실 함수와 최적화 알고리즘이 필요합니다. 일반적으로 다중 클래스 분류 문제에서는 CrossEntropyLoss를 사용하며, 옵티마이저는 SGD나 Adam을 사용하는 경우가 많습니다.

아래의 예제로 함께 살펴보겠습니다.

model = SimpleCNN()
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

for epoch in range(5):
    running_loss = 0.0
    for images, labels in trainloader:
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    print(f"Epoch {epoch+1}, Loss: {running_loss/len(trainloader)}")

위 학습 루프에서는 총 5번의 에폭을 돌며 모델을 학습시킵니다. 학습 중 손실 값이 점차 줄어드는 것을 통해 모델이 데이터를 점점 잘 분류하고 있다는 것을 확인할 수 있습니다.

3. 모델 평가 및 성능 확인

모델 학습이 완료되면, 테스트 데이터셋을 통해 성능을 평가해보는 것이 중요한 포인트 입니다. 이 과정을 통해서 과적합 여부를 판단하거나, 실제 사용 시 어떤 정확도를 기대할 수 있는지를 알 수 있습니다. 평가 시에는 일반적으로 정확도(accuracy)를 사용하며, 혼동 행렬(confusion matrix)이나 F1 점수 등 다양한 지표를 활용할 수 있게 됩니다.

모델 평가 및 성능 확인 코드에 대해서 아래의 예제로 함께 살펴보겠습니다.

correct = 0
total = 0

with torch.no_grad():
    for images, labels in testloader:
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f"정확도: {100 * correct / total:.2f}%")

이 코드는 모델이 테스트 데이터에 대해 얼마나 정확하게 예측했는지를 계산합니다. 실제 프로젝트에서는 단순히 정확도뿐 아니라, 클래스별 정밀도와 재현율 등을 함께 살펴보아야 모델의 편향 여부나 성능을 더 깊이 이해할 수 있습니다.

만약 더 나은 성능을 원한다면 데이터 증강(data augmentation), 드롭아웃(dropout), 배치 정규화(batch normalization) 등의 기법을 모델에 추가적으로 적용해볼 수 있습니다. 또한 학습률 조절이나 더 깊은 네트워크 구조를 사용하여 성능을 향상시킬 수 있습니다. 이러한 과정은 반복적인 실험과 검증을 통해 점차 개선되어야 합니다.

지금까지 PyTorch를 활용한 이미지 분류 모델 구축 과정을 살펴보았습니다. 데이터셋 로딩부터 모델 정의, 학습, 평가까지의 흐름은 딥러닝 프로젝트의 기본적인 구조를 이해하는 데 중요한 기반이 됩니다. 따라서 해당 예제 및 개념이, 실전에 활용할 수 있도록 학습 및 이해하는 데 도움이 되었으면 좋겠습니다.