데이터 준비

PyTorch의 torchvision.datasets 모듈을 사용하여 MNIST 데이터셋을 다운로드하고 불러옵니다. torchvision.transforms.ToTensor()는 이미지를 PyTorch 텐서로 변환하며, torch.utils.data.DataLoader는 데이터를 배치 단위로 효율적으로 로드합니다.

import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np

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

# 데이터셋 다운로드 및 텐서로 변환
# 'data' 폴더에 MNIST 데이터셋 저장
# train=True: 학습용 데이터셋, train=False: 테스트용 데이터셋
# download=True: 데이터가 없으면 다운로드
transform = transforms.Compose([transforms.ToTensor()])
train_dataset = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_dataset = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transform)

# DataLoader 설정
batch_size = 64  # 배치 크기
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)

# 데이터 확인 (선택 사항)
# 첫 번째 배치 데이터 로드
dataiter = iter(train_loader)
images, labels = next(dataiter)

print(f"Image batch shape: {images.shape}") # (batch_size, channels, height, width)
print(f"Label batch shape: {labels.shape}") # (batch_size)

# 데이터 시각화 (선택 사항)
def imshow(img_tensor):
    img = img_tensor.numpy().squeeze() # 텐서를 NumPy 배열로 변환하고 차원 축소
    plt.imshow(img, cmap='gray')
    plt.axis('off')
    plt.show()

# 첫 번째 이미지 시각화
print("Displaying the first image:")
imshow(images[0])
        

CNN 모델 정의

이미지 분류를 위한 간단한 합성곱 신경망(CNN)을 정의합니다. 모델은 합성곱 계층(Conv2d), 최대 풀링 계층(MaxPool2d), ReLU 활성화 함수, 그리고 완전 연결 계층(Linear)으로 구성됩니다. nn.Sequential을 사용하여 모델 구조를 간결하게 정의할 수 있습니다.

# 모델 클래스 정의
class SimpleCNN(nn.Module):
    def __init__(self, num_classes=10):
        super(SimpleCNN, self).__init__()
        # 첫 번째 합성곱 블록
        # 입력 채널 1 (흑백 이미지), 출력 채널 32, 커널 크기 5x5
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=5, padding=2)
        self.relu1 = nn.ReLU()
        # 2x2 최대 풀링
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)

        # 두 번째 합성곱 블록
        # 입력 채널 32, 출력 채널 64, 커널 크기 5x5
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=5, padding=2)
        self.relu2 = nn.ReLU()
        # 2x2 최대 풀링
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)

        # 완전 연결 계층
        # 이전 레이어의 출력 크기 계산 필요: (28x28 -> pool1: 14x14 -> pool2: 7x7)
        # 마지막 풀링 후 채널은 64, 따라서 64 * 7 * 7 = 3136
        self.fc = nn.Linear(in_features=64 * 7 * 7, out_features=num_classes)

    def forward(self, x):
        # 합성곱, ReLU, 풀링 적용
        x = self.pool1(self.relu1(self.conv1(x)))
        x = self.pool2(self.relu2(self.conv2(x)))

        # 출력을 1차원으로 펼침 (batch_size는 유지)
        x = x.view(x.size(0), -1) # Flatten

        # 완전 연결 계층 통과
        x = self.fc(x)
        return x

# 모델 인스턴스 생성 및 GPU로 이동
model = SimpleCNN(num_classes=10).to(device)

# 모델 구조 확인 (선택 사항, torchsummary 라이브러리 필요)
# pip install torchsummary
# from torchsummary import summary
# summary(model, input_size=(1, 28, 28))
print("Model structure created.")
        

모델 학습

손실 함수(nn.CrossEntropyLoss)와 옵티마이저(torch.optim.Adam)를 설정하고, 학습 루프를 실행하여 모델을 훈련합니다. 학습 과정에서는 순전파(forward pass)와 역전파(backward pass)를 반복하며 모델 파라미터를 업데이트합니다. model.train()과 model.eval()은 각각 학습 모드와 평가 모드를 설정합니다.

# 손실 함수와 옵티마이저 정의
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001) # Adam 옵티마이저 사용, 학습률 0.001

# 학습 및 테스트 함수 정의
def train_model(model, device, train_loader, optimizer, criterion, num_epochs):
    model.train() # 모델을 학습 모드로 설정
    for epoch in range(num_epochs):
        running_loss = 0.0
        correct_predictions = 0
        total_samples = 0

        for i, (images, labels) in enumerate(train_loader):
            # 데이터와 레이블을 지정된 장치(GPU/CPU)로 이동
            images, labels = images.to(device), labels.to(device)

            # 옵티마이저의 기울기 초기화
            optimizer.zero_grad()

            # 순전파: 모델 예측 수행
            outputs = model(images)

            # 손실 계산
            loss = criterion(outputs, labels)

            # 역전파: 손실로부터 기울기 계산
            loss.backward()

            # 옵티마이저: 모델 파라미터 업데이트
            optimizer.step()

            # 통계 계산
            running_loss += loss.item() * images.size(0)
            _, predicted = torch.max(outputs.data, 1)
            total_samples += labels.size(0)
            correct_predictions += (predicted == labels).sum().item()

            # 주기적으로 진행 상황 출력
            if (i+1) % 100 == 0:
                print(f'Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{len(train_loader)}], Loss: {loss.item():.4f}')

        epoch_loss = running_loss / total_samples
        epoch_acc = correct_predictions / total_samples
        print(f'--- Epoch {epoch+1} Summary --- Loss: {epoch_loss:.4f}, Accuracy: {epoch_acc*100:.2f}% ---')

    print('Finished Training')

def evaluate_model(model, device, test_loader, criterion):
    model.eval() # 모델을 평가 모드로 설정
    with torch.no_grad(): # 기울기 계산 비활성화
        correct = 0
        total = 0
        test_loss = 0.0
        for images, labels in test_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            loss = criterion(outputs, labels)
            test_loss += loss.item() * images.size(0)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

        avg_test_loss = test_loss / total
        accuracy = correct / total
        print(f'Test Accuracy: {accuracy*100:.2f}%, Average Test Loss: {avg_test_loss:.4f}')
        return accuracy, avg_test_loss

# 학습 파라미터 설정
num_epochs = 5 # 학습 에포크 수

# 모델 학습 실행
train_model(model, device, train_loader, optimizer, criterion, num_epochs)

# 모델 평가 실행
test_accuracy, test_loss_value = evaluate_model(model, device, test_loader, criterion)
        

결과 시각화

학습 과정에서 기록된 훈련 및 테스트 정확도와 손실 값을 그래프로 시각화하여 모델의 성능 변화 추이를 살펴봅니다. 이를 통해 과적합(overfitting) 또는 과소적합(underfitting) 등의 문제를 파악할 수 있습니다.

# 결과 시각화를 위한 데이터 저장 (위 학습 코드에 추가)
train_losses = []
train_accuracies = []
test_losses = []
test_accuracies = []

# 수정된 학습 함수 (결과 저장을 위해)
def train_model_with_logging(model, device, train_loader, optimizer, criterion, num_epochs, test_loader):
    model.train()
    for epoch in range(num_epochs):
        running_loss = 0.0
        correct_predictions = 0
        total_samples = 0

        for i, (images, labels) in enumerate(train_loader):
            images, labels = images.to(device), labels.to(device)
            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            running_loss += loss.item() * images.size(0)
            _, predicted = torch.max(outputs.data, 1)
            total_samples += labels.size(0)
            correct_predictions += (predicted == labels).sum().item()

        epoch_loss = running_loss / total_samples
        epoch_acc = correct_predictions / total_samples
        train_losses.append(epoch_loss)
        train_accuracies.append(epoch_acc)
        print(f'Epoch [{epoch+1}/{num_epochs}], Train Loss: {epoch_loss:.4f}, Train Acc: {epoch_acc*100:.2f}%')

        # 각 에포크마다 테스트 정확도 및 손실 기록
        test_acc, test_loss_val = evaluate_model(model, device, test_loader, criterion)
        test_accuracies.append(test_acc)
        test_losses.append(test_loss_val)

    print('Finished Training')

# 모델 재학습 및 로깅
# num_epochs = 5 # 이전과 동일하게 설정
# train_model_with_logging(model, device, train_loader, optimizer, criterion, num_epochs, test_loader)

# 시각화
plt.figure(figsize=(12, 5))

# 정확도 그래프
plt.subplot(1, 2, 1)
plt.plot(range(1, num_epochs + 1), train_accuracies, label='Train Accuracy')
plt.plot(range(1, num_epochs + 1), test_accuracies, label='Test Accuracy')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.title('Accuracy over Epochs')
plt.legend()

# 손실 그래프
plt.subplot(1, 2, 2)
plt.plot(range(1, num_epochs + 1), train_losses, label='Train Loss')
plt.plot(range(1, num_epochs + 1), test_losses, label='Test Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('Loss over Epochs')
plt.legend()

plt.tight_layout()
plt.show()