PyTorch로 시작하는 CNN 이미지 분류: 기초부터 실전 예제까지

이미지 인식과 처리에 있어 혁신을 가져온 CNN(Convolutional Neural Network)은 현대 딥러닝의 핵심 기술입니다. 많은 개발자들이 딥러닝 입문 시 가장 먼저 접하는 모델이기도 하지만, 실제 PyTorch로 구현할 때 데이터의 차원(Dimension) 처리나 계층(Layer) 설계에서 어려움을 겪곤 합니다.

본 포스팅에서는 PyTorch를 활용하여 CNN 모델을 설계하고, CIFAR-10 데이터셋을 통해 직접 학습 및 평가하는 전 과정을 상세히 다룹니다. 입문자분들이 실무에 바로 적용할 수 있는 표준 코드를 제시합니다.

Generated by Gemini AI.

📌 핵심 요약 3줄

1. CNN의 핵심 구조: 합성곱(Conv), 활성화(ReLU), 풀링(Pooling), 완전연결(FC) 계층의 유기적 결합 이해.
2. PyTorch 구현 실전: nn.Module을 상속받아 유연하고 직관적인 신경망 모델 설계 방법 학습.
3. 모델 최적화 팁: GPU 가속 설정 및 데이터 정규화를 통한 학습 효율 극대화 전략 제공.

---

1. CNN의 기본 개념 및 구성 요소

CNN은 이미지의 공간적 정보를 보존하며 특징을 추출하는 데 최적화되어 있습니다. 주요 구성 요소는 다음과 같습니다.

구성 요소	주요 역할	특징
Convolution Layer	특징 추출 (Feature Extraction)	필터(Kernel)를 통해 이미지의 국소적 패턴 감지
Activation (ReLU)	비선형성 추가	학습의 복잡도를 높여 깊은 신경망 구축 가능
Pooling Layer	차원 축소 (Down-sampling)	계산량 감소 및 주요 특징 강조 (과적합 방지)
Fully Connected Layer	최종 분류 (Classification)	추출된 특징을 1차원 벡터로 변환하여 클래스 결정

---

2. PyTorch로 CNN 모델 구현하기

2.1 라이브러리 임포트 및 환경 설정

가장 먼저 필요한 라이브러리를 불러옵니다. GPU 사용이 가능하다면 device 설정을 추가하는 것이 성능 면에서 유리합니다.

Python

 

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

# GPU 가속 확인
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

2.2 데이터셋 로딩 및 전처리

CIFAR-10 데이터셋은 $32 \times 32$ 크기의 컬러 이미지 60,000장으로 구성되어 있습니다.

Python

 

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 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)

2.3 CNN 모델 정의

nn.Module을 상속받아 모델을 정의합니다. 특히 forward 함수에서 텐서의 형태(Shape)가 어떻게 변화하는지 파악하는 것이 중요합니다.

Python

 

class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        # 입력: (3, 32, 32) -> 출력: (16, 32, 32)
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, padding=1)
        # 입력: (16, 16, 16) -> 출력: (32, 16, 16) - pooling 후 기준
        self.conv2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        
        # MaxPool 2번 적용 시 크기: 32 -> 16 -> 8
        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) # Flatten
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

model = SimpleCNN().to(device)

---

3. 모델 학습 및 평가

3.1 학습 루프 실행

손실 함수는 다중 클래스 분류에 적합한 CrossEntropyLoss를 사용합니다.

Python

 

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

for epoch in range(10):
    running_loss = 0.0
    for inputs, labels in trainloader:
        inputs, labels = inputs.to(device), labels.to(device) # 데이터를 GPU로 이동
        
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        
        running_loss += loss.item()
    
    print(f"Epoch {epoch+1}, Loss: {running_loss/len(trainloader):.4f}")

3.2 모델 성능 평가

학습 시 사용하지 않은 테스트 데이터를 통해 정확도를 측정합니다.

Python

 

correct, total = 0, 0
model.eval() # 평가 모드 전환
with torch.no_grad(): # 기울기 계산 비활성화
    for inputs, labels in testloader:
        inputs, labels = inputs.to(device), labels.to(device)
        outputs = model(inputs)
        _, predicted = torch.max(outputs, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f'Accuracy on test images: {100 * correct / total:.2f}%')

---

💡 개발 꿀팁 및 흔히 하는 실수

구분	내용
꿀팁: GPU 활용	반드시 model.to(device)와 data.to(device)를 모두 호출해야 연산 에러가 발생하지 않습니다.
꿀팁: 평가 모드	평가 시 model.eval()을 호출하면 Dropout이나 Batch Normalization이 평가용으로 동작합니다.
흔한 실수: 차원 불일치	nn.Linear의 입력 차원 계산 실수(32 * 8 * 8)가 잦습니다. 중간 과정의 텐서 크기를 print(x.shape)로 찍어보세요.
흔한 실수: Zero Grad	루프마다 optimizer.zero_grad()를 빼먹으면 이전 기울기가 누적되어 학습이 망가집니다.

---

4. 결론

이번 포스팅에서는 PyTorch를 이용해 가장 기초적인 CNN의 구조를 설계하고 구현해 보았습니다. CIFAR-10과 같은 소규모 데이터셋으로 기본기를 다진 후, 더 깊은 네트워크(ResNet, VGG 등)나 데이터 증강(Data Augmentation) 기법을 적용해 성능을 고도화해 보시기 바랍니다.

PyTorch의 유연한 구조는 여러분의 아이디어를 코드로 구현하는 데 강력한 도구가 될 것입니다. 질문이 있으시면 댓글로 남겨주세요!