PyTorch를 활용한 MNIST 손글씨 숫자 인식 모델 구축

이 문서에서는 PyTorch 라이브러리를 사용하여 MNIST 손글씨 숫자 데이터셋을 분류하는 간단한 합성곱 신경망(CNN) 모델을 구현하는 과정을 다룹니다. 데이터 준비부터 모델 정의, 훈련, 평가 및 결과 시각화까지 전반적인 파이프라인을 설명합니다.

1. 환경 설정 및 데이터 준비

모델 훈련을 위한 환경을 설정하고 필요한 데이터를 로드합니다. GPU 사용 가능 여부를 확인하여 연산을 가속화하고, PyTorch의 torchvision 모듈을 통해 MNIST 데이터셋을 다운로드 및 준비합니다.

import torch
import torch.nn as nn
import matplotlib.pyplot as plt
import torchvision
from torch.utils.data import DataLoader, Dataset
import numpy as np
import torch.nn.functional as F

# GPU 사용 가능 여부 확인 및 장치 설정
compute_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"현재 학습에 사용될 장치: {compute_device}")

MNIST 데이터셋은 손글씨 숫자 이미지와 해당 레이블로 구성되어 있습니다. torchvision.datasets.MNIST를 사용하여 훈련 세트와 테스트 세트를 로드하며, 이미지 데이터를 PyTorch 텐서 형식으로 변환합니다.

# MNIST 데이터셋 다운로드 및 로드
# transform.ToTensor()는 PIL Image를 PyTorch Tensor로 변환하고 픽셀 값을 [0, 1] 범위로 정규화합니다.
train_dataset = torchvision.datasets.MNIST('data', 
                                          train=True, 
                                          transform=torchvision.transforms.ToTensor(), 
                                          download=True)

test_dataset = torchvision.datasets.MNIST('data', 
                                         train=False, 
                                         transform=torchvision.transforms.ToTensor(), 
                                         download=True)

Dataset과 DataLoader 이해

Dataset은 샘플과 해당 레이블을 저장하며, __len____getitem__ 두 가지 함수를 필수적으로 구현해야 합니다. __len__은 데이터셋의 총 샘플 수를 반환하고, __getitem__은 주어진 인덱스에 해당하는 샘플을 반환합니다.

DataLoaderDataset을 감싸 배치(batch) 단위로 데이터를 효율적으로 불러올 수 있도록 도와주는 이터러블 객체입니다. 훈련 중 데이터 셔플링, 병렬 데이터 로딩 등 다양한 기능을 제공합니다.

  • dataset: 로드할 Dataset 객체.
  • batch_size: 각 배치에 포함될 샘플의 수.
  • shuffle: 각 에포크마다 데이터를 섞을지 여부. 훈련 시에는 일반적으로 True로 설정합니다.
  • num_workers: 데이터를 로드할 때 사용할 서브 프로세스 수. 0은 메인 프로세스에서 로드함을 의미합니다.
batch_size = 64 # 배치 크기 설정

# DataLoader 생성
train_loader = DataLoader(train_dataset, 
                          batch_size=batch_size, 
                          shuffle=True)

test_loader = DataLoader(test_dataset, 
                         batch_size=batch_size)

데이터 로더에서 첫 번째 배치를 가져와 데이터의 형태(shape)를 확인합니다. MNIST 이미지는 단일 채널(흑백), 28x28 픽셀 크기를 가집니다.

# 첫 번째 배치 데이터를 가져와 형태 확인
sample_images, sample_labels = next(iter(train_loader))
print(f"이미지 배치 형태: {sample_images.shape}")
print(f"레이블 배치 형태: {sample_labels.shape}")
이미지 배치 형태: torch.Size([64, 1, 28, 28])
레이블 배치 형태: torch.Size([64])

데이터 시각화를 통해 MNIST 이미지들이 어떻게 구성되어 있는지 확인합니다.

# 일부 이미지 시각화
plt.figure(figsize=(15, 4))
for i, img_data in enumerate(sample_images[:10]):
    np_img = np.squeeze(img_data.numpy()) # 차원 축소 (1, 28, 28) -> (28, 28)
    plt.subplot(2, 5, i + 1)
    plt.imshow(np_img, cmap=plt.cm.binary)
    plt.title(f"레이블: {sample_labels[i].item()}")
    plt.axis('off')
plt.suptitle("MNIST 데이터셋 샘플 이미지", fontsize=16)
plt.show()
MNIST Sample Images

2. CNN 모델 정의

손글씨 숫자 분류를 위한 간단한 합성곱 신경망(Convolutional Neural Network, CNN) 모델을 정의합니다. 일반적인 CNN은 이미지 특징을 추출하는 합성곱 레이어와 분류를 위한 완전 연결 레이어로 구성됩니다.

class SimpleCNNClassifier(nn.Module):
    def __init__(self, num_output_classes=10):
        super().__init__()
        # 특징 추출을 위한 합성곱 블록
        self.feature_extractor = nn.Sequential(
            nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3, padding=0), # 첫 번째 합성곱 레이어
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2), # 첫 번째 풀링 레이어

            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding=0), # 두 번째 합성곱 레이어
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)  # 두 번째 풀링 레이어
        )
        
        # 분류를 위한 완전 연결 레이어
        # 28x28 -> Conv1(k=3,p=0) -> 26x26 -> Pool1(k=2) -> 13x13
        # 13x13 -> Conv2(k=3,p=0) -> 11x11 -> Pool2(k=2) -> 5x5
        # 최종 특징 맵 크기: 64채널 * 5 * 5 = 1600
        self.classifier_head = nn.Sequential(
            nn.Linear(64 * 5 * 5, 128), # 첫 번째 완전 연결 레이어 (출력 채널을 128로 변경)
            nn.ReLU(),
            nn.Linear(128, num_output_classes) # 최종 출력 레이어
        )

    def forward(self, input_data):
        features = self.feature_extractor(input_data)
        # 완전 연결 레이어에 전달하기 위해 텐서를 평탄화합니다.
        features = torch.flatten(features, 1) # 배치 차원 제외하고 평탄화
        logits = self.classifier_head(features)
        return logits

# 모델 인스턴스화 및 GPU로 이동
model_instance = SimpleCNNClassifier(num_output_classes=10).to(compute_device)

# 모델 구조 요약
try:
    from torchinfo import summary
    summary(model_instance, input_size=(1, 1, 28, 28))
except ImportError:
    print("torchinfo가 설치되지 않았습니다. pip install torchinfo를 실행하여 모델 요약을 확인하세요.")
    print(model_instance)
현재 학습에 사용될 장치: cuda
==========================================================================================
Layer (type:depth-idx)                   Output Shape              Param #
==========================================================================================
SimpleCNNClassifier                      [1, 10]                   --
├─Sequential: 1-1                        [1, 64, 5, 5]             --
│    ├─Conv2d: 2-1                       [1, 32, 26, 26]           320
│    ├─ReLU: 2-2                         [1, 32, 26, 26]           --
│    ├─MaxPool2d: 2-3                    [1, 32, 13, 13]           --
│    ├─Conv2d: 2-4                       [1, 64, 11, 11]           18,496
│    ├─ReLU: 2-5                         [1, 64, 11, 11]           --
│    └─MaxPool2d: 2-6                    [1, 64, 5, 5]             --
├─Sequential: 1-2                        [1, 10]                   --
│    ├─Linear: 2-7                       [1, 128]                  204,928
│    ├─ReLU: 2-8                         [1, 128]                  --
│    └─Linear: 2-9                       [1, 10]                   1,290
==========================================================================================
Total params: 225,034
Trainable params: 225,034
Non-trainable params: 0
==========================================================================================

모델의 각 구성 요소는 다음과 같은 역할을 합니다:

  • nn.Conv2d: 이미지에서 특징(feature)을 추출하는 합성곱 레이어. 입력 채널, 출력 채널, 필터(커널) 크기 등을 정의합니다.
  • nn.MaxPool2d: 특징 맵의 크기를 줄여(다운샘플링) 계산 비용을 줄이고 과적합을 방지하는 풀링 레이어.
  • nn.ReLU: 비선형성을 도입하여 모델이 복잡한 패턴을 학습할 수 있도록 돕는 활성화 함수.
  • nn.Linear: 완전 연결 레이어로, 특징 추출 후 분류를 수행합니다. 입력 특징 수와 출력 특징 수를 정의합니다.

3. 모델 훈련 설정

모델 훈련을 위한 손실 함수와 최적화 알고리즘을 설정합니다. 다중 클래스 분류 문제이므로 nn.CrossEntropyLoss를 사용하고, 옵티마이저는 SGD(Stochastic Gradient Descent)를 선택합니다.

# 손실 함수 및 최적화기 설정
loss_criterion = nn.CrossEntropyLoss() # 크로스 엔트로피 손실 함수
learning_rate = 0.01                 # 학습률
optimizer = torch.optim.SGD(model_instance.parameters(), lr=learning_rate) # SGD 옵티마이저

훈련 및 테스트 과정을 담당하는 함수를 각각 정의합니다.

# 모델 훈련 함수
def train_epoch(data_loader, model, loss_func, optimizer_algo):
    model.train() # 모델을 훈련 모드로 설정
    total_samples = len(data_loader.dataset)
    num_batches = len(data_loader)
    
    current_loss_sum, correct_predictions = 0, 0
    
    for inputs, targets in data_loader:
        inputs, targets = inputs.to(compute_device), targets.to(compute_device)
        
        # 예측 및 손실 계산
        predictions = model(inputs)
        loss = loss_func(predictions, targets)
        
        # 역전파 및 가중치 업데이트
        optimizer_algo.zero_grad() # 이전 단계의 기울기 초기화
        loss.backward()            # 손실에 대한 기울기 계산
        optimizer_algo.step()      # 옵티마이저를 사용하여 모델 파라미터 업데이트
        
        # 정확도 및 손실 누적
        correct_predictions += (predictions.argmax(1) == targets).type(torch.float).sum().item()
        current_loss_sum += loss.item()
            
    avg_accuracy = correct_predictions / total_samples
    avg_loss = current_loss_sum / num_batches

    return avg_accuracy, avg_loss

# 모델 테스트/평가 함수
def evaluate_model(data_loader, model, loss_func):
    model.eval() # 모델을 평가 모드로 설정
    total_samples = len(data_loader.dataset)
    num_batches = len(data_loader)
    
    current_loss_sum, correct_predictions = 0, 0
    
    # 기울기 계산 비활성화 (메모리 절약 및 연산 속도 향상)
    with torch.no_grad():
        for inputs, targets in data_loader:
            inputs, targets = inputs.to(compute_device), targets.to(compute_device)
            
            predictions = model(inputs)
            loss = loss_func(predictions, targets)
            
            current_loss_sum += loss.item()
            correct_predictions += (predictions.argmax(1) == targets).type(torch.float).sum().item()

    avg_accuracy = correct_predictions / total_samples
    avg_loss = current_loss_sum / num_batches

    return avg_accuracy, avg_loss

위 함수에서 사용된 핵심 옵티마이저 메서드에 대한 설명은 다음과 같습니다:

  • optimizer.zero_grad(): 이전 미니배치에서 계산된 모든 기울기(gradient)를 0으로 초기화합니다. 이는 기울기가 누적되는 것을 방지하기 위함입니다.
  • loss.backward(): 현재 손실(loss)에 대해 역전파를 수행하고, 각 학습 가능한 파라미터에 대한 기울기를 계산하여 해당 파라미터의 .grad 속성에 저장합니다.
  • optimizer.step(): 계산된 기울기를 사용하여 모델의 파라미터를 업데이트합니다. 이는 옵티마이저에 정의된 규칙(예: SGD, Adam)에 따라 파라미터 값을 조정합니다.

또한, model.train()model.eval()은 모델의 동작 모드를 변경합니다. 드롭아웃(Dropout)이나 배치 정규화(Batch Normalization)와 같은 레이어가 포함된 모델의 경우, 이 모드 설정은 올바른 동작을 위해 필수적입니다. model.train()은 이러한 레이어들을 활성화하여 훈련 중 예측 불확실성을 높이고 일반화 성능을 향상시키는 반면, model.eval()은 이들을 비활성화하여 일관된 평가를 보장합니다.

4. 모델 훈련 실행

정의된 훈련 및 테스트 함수를 사용하여 모델을 여러 에포크(epoch) 동안 훈련시킵니다. 각 에포크마다 훈련 및 테스트 세트의 정확도와 손실을 기록하고 출력합니다.

num_epochs = 10 # 훈련 에포크 수

# 훈련 및 테스트 기록을 위한 리스트
history_train_loss = []
history_train_acc = []
history_test_loss = []
history_test_acc = []

print("모델 훈련 시작...")
for epoch_idx in range(num_epochs):
    # 훈련 단계
    epoch_train_acc, epoch_train_loss = train_epoch(train_loader, model_instance, loss_criterion, optimizer)
    
    # 평가 단계
    epoch_test_acc, epoch_test_loss = evaluate_model(test_loader, model_instance, loss_criterion)
    
    # 결과 기록
    history_train_acc.append(epoch_train_acc)
    history_train_loss.append(epoch_train_loss)
    history_test_acc.append(epoch_test_acc)
    history_test_loss.append(epoch_test_loss)
    
    # 에포크 결과 출력
    log_message = (f"에포크 {epoch_idx+1:2d}: "
                   f"훈련 정확도: {epoch_train_acc*100:.2f}%, "
                   f"훈련 손실: {epoch_train_loss:.4f}, "
                   f"테스트 정확도: {epoch_test_acc*100:.2f}%, "
                   f"테스트 손실: {epoch_test_loss:.4f}")
    print(log_message)
print("훈련 완료!")
모델 훈련 시작...
에포크  1: 훈련 정확도: 91.67%, 훈련 손실: 0.2831, 테스트 정확도: 97.46%, 테스트 손실: 0.0890
에포크  2: 훈련 정확도: 97.64%, 훈련 손실: 0.0784, 테스트 정확도: 98.12%, 테스트 손실: 0.0614
에포크  3: 훈련 정확도: 98.42%, 훈련 손실: 0.0526, 테스트 정확도: 98.56%, 테스트 손실: 0.0461
에포크  4: 訓練 정확도: 98.81%, 훈련 손실: 0.0388, 테스트 정확도: 98.71%, 테스트 손실: 0.0392
에포크  5: 訓練 정확도: 99.07%, 훈련 손실: 0.0305, 테스트 정확도: 98.84%, 테스트 손실: 0.0347
에포크  6: 訓練 정확도: 99.27%, 훈련 손실: 0.0242, 테스트 정확도: 98.88%, 테스트 손실: 0.0306
에포크  7: 訓練 정확도: 99.41%, 훈련 손실: 0.0194, 테스트 정확도: 99.04%, 테스트 손실: 0.0287
에포크  8: 訓練 정확도: 99.52%, 훈련 손실: 0.0160, 테스트 정확도: 99.07%, 테스트 손실: 0.0267
에포크  9: 訓練 정확도: 99.64%, 훈련 손실: 0.0125, 테스트 정확도: 99.11%, 테스트 손실: 0.0267
에포크 10: 訓練 정확도: 99.71%, 훈련 손실: 0.0102, 테스트 정확도: 99.04%, 테스트 손실: 0.0270
훈련 완료!

5. 훈련 결과 시각화

훈련 과정 동안 기록된 정확도와 손실 값을 그래프로 시각화하여 모델의 학습 추이를 분석합니다.

# 시각화를 위한 설정
# 한글 폰트 설정 및 음수 부호 처리 (matplotlib에서 한글 깨짐 방지)
plt.rcParams['font.family'] = 'Malgun Gothic' # Windows 기준, macOS는 'AppleGothic'
plt.rcParams['axes.unicode_minus'] = False
plt.rcParams['figure.dpi'] = 100 # 해상도

epoch_indices = range(num_epochs)

plt.figure(figsize=(12, 4))

# 정확도 그래프
plt.subplot(1, 2, 1)
plt.plot(epoch_indices, history_train_acc, label='훈련 정확도')
plt.plot(epoch_indices, history_test_acc, label='테스트 정확도')
plt.legend(loc='lower right')
plt.title('훈련 및 테스트 정확도')
plt.xlabel('에포크')
plt.ylabel('정확도')

# 손실 그래프
plt.subplot(1, 2, 2)
plt.plot(epoch_indices, history_train_loss, label='훈련 손실')
plt.plot(epoch_indices, history_test_loss, label='테스트 손실')
plt.legend(loc='upper right')
plt.title('훈련 및 테스트 손실')
plt.xlabel('에포크')
plt.ylabel('손실')

plt.tight_layout() # 그래프 간 간격 자동 조절
plt.show()
Training and Validation Accuracy and Loss Plots

태그: PyTorch mnist CNN 손글씨 인식 딥러닝

8월 7일 05:42에 게시됨