MNIST 학습
PyTorch에서 MNIST 데이터셋을 학습하는 과정은 딥러닝의 기본적인 워크플로우를 이해하는 좋은 시작점입니다. MNIST는 손으로 쓴 숫자 이미지(28x28 픽셀, 흑백)를 분류하는 작업으로 구성되어 있으며, 클래스는 0부터 9까지 총 10개입니다.
아래에 MNIST 데이터셋을 활용하여 PyTorch로 딥러닝 모델을 구현하는 전체 과정을 단계별로 설명하겠습니다.
1. PyTorch의 기본 구성
PyTorch에서 MNIST 작업은 아래와 같은 주요 단계로 구성됩니다: 1. 데이터 준비 2. 모델 정의 3. 학습 설정 4. 학습 실행 5. 평가
2. 단계별 코드 설명
2.1. 라이브러리 임포트
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
2.2. 데이터 준비
torchvision.datasets를 사용하여 MNIST 데이터셋을 다운로드하고 변환합니다.
# 데이터 변환: Tensor 변환과 정규화
transform = transforms.Compose([
transforms.ToTensor(), # 이미지를 Tensor로 변환
transforms.Normalize((0.5,), (0.5,)) # 평균 0.5, 표준편차 0.5로 정규화
])
# 학습 데이터셋과 테스트 데이터셋 로드
train_dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)
test_dataset = datasets.MNIST(root='./data', train=False, transform=transform, download=True)
# 데이터 로더 생성
train_loader = DataLoader(dataset=train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(dataset=test_dataset, batch_size=64, shuffle=False)
2.3. 모델 정의
신경망 모델은 PyTorch의 nn.Module을 상속받아 정의합니다. 여기서는 간단한 Fully Connected Neural Network를 사용합니다.
class MNISTModel(nn.Module):
def __init__(self):
super(MNISTModel, self).__init__()
self.fc = nn.Sequential(
nn.Flatten(), # 28x28 이미지를 1차원 벡터로 변환
nn.Linear(28*28, 128), # 입력 크기: 784, 출력 크기: 128
nn.ReLU(), # 활성화 함수
nn.Linear(128, 64), # 128 -> 64
nn.ReLU(),
nn.Linear(64, 10) # 64 -> 10 (클래스 개수)
)
def forward(self, x):
return self.fc(x)
# 모델 초기화
model = MNISTModel()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device) # GPU로 이동
2.4. 손실 함수 및 최적화 설정
- 손실 함수: 분류 문제에서는 CrossEntropyLoss 사용.
- 최적화 기법: Adam 옵티마이저 사용.
criterion = nn.CrossEntropyLoss() # 손실 함수
optimizer = optim.Adam(model.parameters(), lr=0.001) # Adam 옵티마이저
2.5. 학습 루프
MNIST 데이터셋을 모델에 학습시킵니다. 각 epoch에서 데이터를 순차적으로 가져와 학습합니다.
epochs = 5
for epoch in range(epochs):
model.train() # 모델을 학습 모드로 전환
running_loss = 0.0
for images, labels in train_loader:
images, labels = images.to(device), labels.to(device) # GPU로 데이터 이동
# Forward pass
outputs = model(images)
loss = criterion(outputs, labels)
# Backward pass and optimization
optimizer.zero_grad() # Gradient 초기화
loss.backward() # Backpropagation
optimizer.step() # Optimizer 업데이트
running_loss += loss.item()
print(f"Epoch [{epoch+1}/{epochs}], Loss: {running_loss/len(train_loader):.4f}")
2.6. 모델 평가
학습된 모델의 성능을 테스트 데이터셋으로 평가합니다.
model.eval() # 평가 모드로 전환
correct = 0
total = 0
with torch.no_grad(): # Gradient 계산 비활성화
for images, labels in test_loader:
images, labels = images.to(device), labels.to(device)
outputs = model(images)
_, predicted = torch.max(outputs, 1) # 가장 높은 확률의 클래스 예측
total += labels.size(0)
correct += (predicted == labels).sum().item()
print(f'Accuracy: {100 * correct / total:.2f}%')
3. 확장 및 최적화
3.1. CNN 적용
Fully Connected Neural Network 대신 Convolutional Neural Network(CNN)를 사용하면 성능이 향상됩니다.
class CNNModel(nn.Module):
def __init__(self):
super(CNNModel, self).__init__()
self.conv_layers = nn.Sequential(
nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),
nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2)
)
self.fc_layers = nn.Sequential(
nn.Flatten(),
nn.Linear(64*7*7, 128),
nn.ReLU(),
nn.Linear(128, 10)
)
def forward(self, x):
x = self.conv_layers(x)
x = self.fc_layers(x)
return x
# CNN 모델 생성
model = CNNModel().to(device)
3.2. Mixed Precision Training
Automatic Mixed Precision(AMP)를 적용하여 GPU 메모리 사용량을 줄이고 연산 속도를 높일 수 있습니다.
4. 결과
위 코드를 실행하면 5 epoch 기준 약 98% 이상의 정확도를 얻을 수 있습니다. MNIST 데이터셋은 PyTorch로 딥러닝을 학습하고, 모델 설계 및 최적화를 실습하기에 적합한 예제입니다. Jetson Orin Nano 같은 엣지 디바이스에서도 효율적으로 실행할 수 있습니다.