2026-01-27 | 8 min read | 150 | PyTorch CNN 이미지분류

6-4. CNN: PyTorch로 CNN 구현

1. CNN 모델 구현 개요

합성곱 신경망(Convolutional Neural Network, CNN)은 이미지 인식, 객체 탐지, 이미지 생성 등 다양한 컴퓨터 비전 분야에서 괄목할 만한 성능을 보이는 딥러닝 모델입니다. 이전 챕터에서 CNN의 기본 개념과 구조를 살펴보았습니다. 이제 PyTorch를 사용하여 실제 CNN 모델을 구현하고, 학습시키며, 이미지 분류를 수행하는 과정을 자세히 알아보겠습니다.

CNN을 구현하는 것은 딥러닝을 이해하는 데 매우 중요한 단계입니다. PyTorch는 텐서 연산, 자동 미분, 모델 구축 등 딥러닝 모델 개발에 필요한 다양한 기능을 제공하며, 직관적인 API를 통해 CNN을 쉽게 구현할 수 있도록 돕습니다.

2. PyTorch 환경 설정

PyTorch를 사용하기 위해서는 먼저 PyTorch를 설치해야 합니다. 설치 방법은 다음과 같습니다.

pip install torch torchvision torchaudio

CUDA를 지원하는 GPU를 사용하고 싶다면, CUDA 버전에 맞는 PyTorch 버전을 설치해야 합니다. PyTorch 공식 웹사이트에서 설치 가이드를 참고하여 적절한 버전을 선택하세요.

3. CNN 모델 정의

PyTorch에서 CNN 모델을 정의하기 위해 torch.nn 모듈을 사용합니다. torch.nn 모듈은 신경망을 구성하는 데 필요한 다양한 레이어를 제공합니다.

1) Convolutional Layer (`nn.Conv2d`)

합성곱 레이어는 CNN의 핵심 구성 요소입니다. 입력 이미지에 필터를 적용하여 특징 맵을 생성합니다. PyTorch에서는 nn.Conv2d 클래스를 사용하여 합성곱 레이어를 정의합니다.

import torch
import torch.nn as nn

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, stride=1, padding=1)
        # in_channels: 입력 채널 수, out_channels: 출력 채널 수, kernel_size: 필터 크기,
        # stride: 스트라이드, padding: 패딩

    def forward(self, x):
        x = self.conv1(x)
        return x

in_channels: 입력 이미지의 채널 수입니다. 컬러 이미지는 3 (RGB), 흑백 이미지는 1입니다.
out_channels: 출력 채널 수, 즉 필터의 개수입니다. 이는 모델이 학습할 특징의 수를 결정합니다.
kernel_size: 필터의 크기입니다. 예를 들어 kernel_size=3은 3x3 필터를 의미합니다.
stride: 필터가 입력 위를 움직이는 간격입니다.
padding: 입력 이미지의 가장자리에 추가되는 픽셀 수입니다. 패딩은 이미지의 크기를 유지하거나, 경계에 있는 정보를 보존하는 데 사용됩니다.

2) Pooling Layer (`nn.MaxPool2d`, `nn.AvgPool2d`)

풀링 레이어는 특징 맵의 크기를 줄이고, 중요한 특징을 강조하며, 과적합을 방지하는 데 기여합니다. PyTorch에서는 nn.MaxPool2d (최대 풀링) 또는 nn.AvgPool2d (평균 풀링) 클래스를 사용하여 풀링 레이어를 정의합니다.

import torch
import torch.nn as nn

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, stride=1, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)  # kernel_size: 풀링 영역 크기, stride: 스트라이드

    def forward(self, x):
        x = self.conv1(x)
        x = self.pool(x)
        return x

kernel_size: 풀링 영역의 크기입니다.
stride: 풀링 영역이 이동하는 간격입니다. 일반적으로 kernel_size와 같은 값을 사용합니다.

3) 활성화 함수

합성곱 레이어와 풀링 레이어 뒤에는 활성화 함수를 적용하여 모델의 비선형성을 부여합니다. 일반적으로 ReLU (Rectified Linear Unit) 함수를 사용합니다.

import torch
import torch.nn as nn

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, stride=1, padding=1)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu(x)
        return x

4) Fully Connected Layer (`nn.Linear`)

CNN의 마지막 부분에서는 완전 연결 레이어(Fully Connected Layer, FC layer)를 사용하여 특징 맵을 분류합니다. 완전 연결 레이어는 입력으로 1차원 벡터를 받으므로, 특징 맵을 펼쳐서(flatten) 1차원 벡터로 변환해야 합니다.

import torch
import torch.nn as nn

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, stride=1, padding=1)
        self.relu = nn.ReLU()
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(in_features=16 * 16 * 16, out_features=10) # 예시: 입력 이미지 크기 32x32, 10개 클래스

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

in_features: 입력 특징의 수입니다. 이전 레이어의 출력 크기에 따라 달라집니다.
out_features: 출력 특징의 수입니다. 분류 문제의 경우 클래스의 수와 같습니다.
x.view(x.size(0), -1): 특징 맵을 1차원 벡터로 펼치는 역할입니다. x.size(0)은 배치 크기를 유지하고, -1은 나머지 차원을 자동으로 계산하도록 합니다.

5) 전체 CNN 모델 예시

위에서 설명한 레이어들을 조합하여 전체 CNN 모델을 구성할 수 있습니다.

import torch
import torch.nn as nn

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, stride=1, padding=1)
        self.relu1 = nn.ReLU()
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, stride=1, padding=1)
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(in_features=32 * 8 * 8, out_features=128)  # Assuming input size is 32x32
        self.relu3 = nn.ReLU()
        self.fc2 = nn.Linear(in_features=128, out_features=10) # 10 classes

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = x.view(x.size(0), -1)  # flatten
        x = self.fc1(x)
        x = self.relu3(x)
        x = self.fc2(x)
        return x

이 예시는 두 개의 합성곱 블록(합성곱 -> ReLU -> 풀링)과 두 개의 완전 연결 레이어로 구성된 CNN 모델입니다. 입력 이미지 크기와 출력 클래스 수에 따라 in_features와 out_features 값을 적절하게 설정해야 합니다.

4. 데이터셋 준비

CNN 모델을 학습시키기 위해서는 데이터셋이 필요합니다. PyTorch는 torchvision.datasets 모듈을 통해 다양한 데이터셋을 제공합니다. 예를 들어, CIFAR-10 데이터셋을 사용할 수 있습니다.

import torch
import torchvision
import torchvision.transforms as transforms

# 데이터 변환 정의
transform = transforms.Compose(
    [transforms.ToTensor(), # 이미지를 Tensor로 변환
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]) # 정규화

# 데이터셋 로드
trainset <mark class="highlight"><strong><u> torchvision.datasets.CIFAR10(root</u></strong></mark>'./data', train=True,
                                        download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,
                                          shuffle=True, num_workers=2)

testset <mark class="highlight"><strong><u> torchvision.datasets.CIFAR10(root</u></strong></mark>'./data', train=False,
                                       download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=4,
                                         shuffle=False, num_workers=2)

transforms.ToTensor(): 이미지를 텐서로 변환합니다. 픽셀 값을 0과 1 사이의 값으로 정규화합니다.
transforms.Normalize(): 텐서의 각 채널에 대해 평균과 표준편차를 사용하여 정규화합니다. 이 예시에서는 각 채널의 평균이 0.5, 표준편차가 0.5가 되도록 정규화합니다.
torchvision.datasets.CIFAR10: CIFAR-10 데이터셋을 다운로드하고 로드합니다. root는 데이터셋을 저장할 경로, train은 훈련 데이터셋을 사용할지 여부, download는 데이터셋을 다운로드할지 여부, transform은 이미지에 적용할 변환을 지정합니다.
torch.utils.data.DataLoader: 데이터셋을 미니 배치로 나누고, 셔플링하고, 병렬로 데이터를 로드하는 기능을 제공합니다. batch_size는 배치 크기, shuffle은 데이터를 섞을지 여부, num_workers는 데이터를 로드하는 데 사용할 워커 프로세스의 수입니다.

5. 모델 학습

데이터셋과 모델이 준비되었다면, 모델을 학습시킬 수 있습니다. 학습 과정은 다음과 같습니다.

모델 초기화: 앞서 정의한 CNN 모델을 인스턴스화합니다.
손실 함수 정의: 모델의 예측과 실제 레이블 간의 차이를 계산하는 손실 함수를 정의합니다. 일반적으로 교차 엔트로피 손실 함수(Cross-Entropy Loss)를 사용합니다.
옵티마이저 정의: 손실을 최소화하기 위해 모델의 가중치를 업데이트하는 옵티마이저를 정의합니다. 일반적으로 Adam 또는 SGD 옵티마이저를 사용합니다.
학습 루프:
- 미니 배치 단위로 데이터를 로드합니다.
- 모델에 입력을 전달하여 예측값을 계산합니다 (forward pass).
- 손실 함수를 사용하여 예측값과 실제 레이블 간의 손실을 계산합니다.
- backward() 함수를 사용하여 손실을 역전파하고, 그래디언트를 계산합니다.
- step() 함수를 사용하여 옵티마이저가 모델의 가중치를 업데이트합니다.

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

# 1. 모델 초기화
model = CNN()

# 2. 손실 함수 정의
criterion = nn.CrossEntropyLoss() # 다중 분류 문제에 적합

# 3. 옵티마이저 정의
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 데이터 변환 정의
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

# 데이터셋 로드
trainset <mark class="highlight"><strong><u> datasets.CIFAR10(root</u></strong></mark>'./data', train=True, download=True, transform=transform)
trainloader = DataLoader(trainset, batch_size=4, shuffle=True, num_workers=2)

# 4. 학습 루프
num_epochs = 2  # 에폭 수
for epoch in range(num_epochs):
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        # 입력 데이터와 레이블 가져오기
        inputs, labels = data

        # 그래디언트 초기화
        optimizer.zero_grad()

        # 순전파 + 역전파 + 최적화
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        # 손실 통계
        running_loss += loss.item()
        if i % 2000 == 1999:    # print every 2000 mini-batches
            print('[%d, %5d] loss: %.3f' %
                  (epoch + 1, i + 1, running_loss / 2000))
            running_loss = 0.0

print('Finished Training')

nn.CrossEntropyLoss(): 교차 엔트로피 손실 함수를 사용합니다. 다중 클래스 분류 문제에 적합하며, 모델의 출력에 softmax 함수를 자동으로 적용합니다.
optim.Adam(): Adam 옵티마이저를 사용합니다. model.parameters()는 모델의 모든 학습 가능한 파라미터를 반환합니다. lr은 학습률을 나타냅니다.
optimizer.zero_grad(): 각 미니 배치마다 그래디언트를 초기화합니다. 이전 미니 배치의 그래디언트는 현재 미니 배치에 영향을 미치지 않도록 합니다.
loss.backward(): 손실을 역전파하여 그래디언트를 계산합니다.
optimizer.step(): 옵티마이저가 모델의 가중치를 업데이트합니다.

학습 루프 설명 뒤

6. 모델 평가 및 이미지 분류

모델 학습이 완료되면, 테스트 데이터셋을 사용하여 모델의 성능을 평가할 수 있습니다.

correct = 0
total = 0
with torch.no_grad(): # 그래디언트 계산 비활성화
    for data in testloader:
        images, labels = data
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1) # 예측 레이블
        total += labels.size(0)
        correct += (predicted <mark class="highlight"> labels).sum().item()

print('Accuracy of the network on the 10000 test images: %d %%' % (
    100 * correct / total))

torch.no_grad(): 평가 과정에서는 그래디언트를 계산할 필요가 없으므로, 메모리 사용량을 줄이고 계산 속도를 높이기 위해 torch.no_grad() 컨텍스트를 사용합니다.
torch.max(outputs.data, 1): 모델의 출력 중 가장 높은 값을 갖는 인덱스를 반환합니다. 이 인덱스는 예측된 클래스를 나타냅니다.
labels.size(0): 배치 크기를 반환합니다.
(predicted </mark> labels).sum().item(): 예측값과 실제 레이블이 일치하는 개수를 계산합니다.

모델의 정확도를 측정한 후, 개별 이미지를 분류하는 데 사용할 수 있습니다.

import matplotlib.pyplot as plt
import numpy as np

# 테스트 이미지 표시 함수
def imshow(img):
    img = img / 2 + 0.5  # unnormalize
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()

# 테스트 이미지 샘플 가져오기
dataiter = iter(testloader)
images, labels = next(dataiter)

# 이미지 예측
outputs = model(images)
_, predicted = torch.max(outputs, 1)

# 이미지 표시
imshow(torchvision.utils.make_grid(images))
print('Predicted: ', ' '.join('%5s' % classes[predicted[j]]
                              for j in range(4)))

matplotlib.pyplot: 이미지를 시각화하는 데 사용됩니다.
torchvision.utils.make_grid: 여러 이미지를 하나의 그리드로 묶는 데 사용됩니다.

7. GPU 사용 (선택 사항)

모델 학습 속도를 향상시키기 위해 GPU를 사용할 수 있습니다. torch.cuda.is_available() 함수를 사용하여 GPU 사용 가능 여부를 확인하고, 모델과 데이터를 GPU로 이동합니다.

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

# 모델을 GPU로 이동
model.to(device)

# 입력 데이터와 레이블을 GPU로 이동
inputs, labels = data[0].to(device), data[1].to(device)

8. 모델 저장 및 로드 (선택 사항)

학습된 모델을 저장하고 나중에 다시 사용할 수 있습니다.

# 모델 저장
torch.save(model.state_dict(), './cnn_model.pth')

# 모델 로드
model = CNN()
model.load_state_dict(torch.load('./cnn_model.pth'))

torch.save(): 모델의 가중치를 저장합니다. state_dict()는 모델의 가중치와 관련된 정보를 딕셔너리 형태로 반환합니다.
model.load_state_dict(): 저장된 가중치를 모델에 로드합니다.

9. 결론

PyTorch를 사용하여 CNN 모델을 구현하고, 학습하고, 이미지 분류를 수행하는 과정을 살펴보았습니다. CNN은 딥러닝에서 매우 중요한 모델이며, PyTorch를 사용하면 CNN을 쉽게 구현하고 실험할 수 있습니다. 본 가이드에서 제시된 예시를 바탕으로 다양한 CNN 구조를 직접 구현해보고, 데이터셋을 변경해보면서 CNN에 대한 이해를 높여보세요.

CNN은 이미지 처리 분야뿐만 아니라 자연어 처리, 음성 인식 등 다양한 분야에서 활용될 수 있습니다. 딥러닝 기술을 발전시키기 위해 CNN을 꾸준히 학습하고, 다양한 응용 분야에 적용하는 연습을 해보세요.