1.CNN으로 MNIST 학습하기
흐름
1. import
2. 데이터 변환 & 로딩 (MNIST)
3. CNN 모델 정의 (Conv2d → ReLU → MaxPool2d → Flatten → Linear)
4. 손실 함수 & 옵티마이저 설정
5. 1 Epoch 학습 루프
6. 출력 shape, loss 확인
코드
# 0. Import
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
# 1. 데이터 변환 및 로딩
transform = transforms.ToTensor()
train_dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
# 2. CNN 모델 정의
class CNN(nn.Module):
def __init__(self): #layer 정의
super().__init__()
self.conv1 = nn.Conv2d( # (1) 2D Convolution: 공간 구조 유지하며 특징 추출
in_channels=1, # MNIST는 흑백이므로 채널 1개
out_channels=16, # 필터(커널) 16개 사용 → 특징맵 16개 생성
kernel_size=3, # 특징(3x3)필터
padding=1 # 가장자리 보존을 위해 zero-padding 1 적용
) # → 입력: [B, 1, 28, 28] → 출력: [B, 16, 28, 28]
self.pool = nn.MaxPool2d(2, 2) # (2) Max Pooling: 특징 유지하며 크기 줄이기
# kernel_size=2 : 2×2 영역을 한 블록으로 묶음
# tride=2 : 2칸씩 이동하며 풀링 → 겹치지 않음
# 연산 방식 : 2×2 안에서 최댓값 하나만 선택
# → [B, 16, 28, 28] → [B, 16, 14, 14]
self.fc1 = nn.Linear(16 * 14 * 14, 10) #최종 분류
# 16 * 14 * 14 : 이전 Conv + Pool 연산으로 얻은 이미지 특징의 총 개수
# 10 : 최종 클래스 수 (MNIST는 숫자 0~9 → 총 10개)
# 역할 : 16 * 16 * 14를 1차원으로 표서 10차원으로 변환(출력 10개 0~9로)
def forward(self, x): # 실행
x = self.pool(F.relu(self.conv1(x))) # 합성곱(특징 추출) + relu(특징 추출 도움) + pool로 단순화
# x.shape == [B, 16, 14, 14]
x = x.view(x.size(0), -1) # 평탄화(Flatten)
# size(0)은 B 즉 배치 크기 가져오고 -1은 나머지 숫자들 전부 곱하라는 뜻임
# [배치 크기, 특징 벡터 차원] → [64, 3136] 결국 이 형태로 됨
x = self.fc1(x) # Fully connected(완전 연결층으로 분류 마무리 0~9)
return x
model = CNN()
# 3. 손실 함수 & 옵티마이저
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
# 4. 학습 루프 (1 Epoch만)
for batch_idx, (images, labels) in enumerate(train_loader):
outputs = model(images)
loss = criterion(outputs, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if batch_idx % 100 == 0:
print(f"Batch {batch_idx}, Loss: {loss.item():.4f}")
2.MNIST를 각각 MLP와 CNN으로 학습시켰을 때의 성능 비교(5 epoch만큼 반복)
MLP로 MNIST를 학습시킨 경우와 CNN으로 MNIST를 학습시킨 경우간의 Loss차이를 비교하여 성능 차이를 확인한다.
(모델 관련 코드를 제외한 다른 코드들은 모두 통일하고, epoch수를 5로 늘려 성능 차이가 잘 드러내도록 한다.)
(1)MLP로 MNIST를 학습시킨 경우
코드
# 0. import
import torch
import torch.nn as nn
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
# 1. 데이터 변환 설정
transform = transforms.ToTensor()
# 2. 데이터셋 불러오기
train_dataset = datasets.MNIST(
root='./data',
train=True,
transform=transform,
download=True
)
# 3. 데이터로더 설정
train_loader = DataLoader(
dataset=train_dataset,
batch_size=64,
shuffle=True
)
# 4. 모델 설정
model = nn.Sequential(
nn.Flatten(), # [B, 1, 28, 28] → [B, 784]
nn.Linear(784, 128), # 은닉층 (128개 노드)
nn.ReLU(), # 활성화 함수
nn.Linear(128, 10) # 출력층 (10개 클래스)
)
# 5. 옵티마이저, 손실 함수 설정
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
criterion = nn.CrossEntropyLoss()
# 6. 학습 루프 (5 Epoch)
num_epochs = 5
for epoch in range(num_epochs):
for batch_idx, (images, labels) in enumerate(train_loader):
outputs = model(images) # 순전파
loss = criterion(outputs, labels) # 손실 계산
optimizer.zero_grad() # 기울기 초기화
loss.backward() # 역전파
optimizer.step() # 가중치 업데이트
if batch_idx % 100 == 0:
print(f"[Epoch {epoch+1}] Batch {batch_idx}, Loss: {loss.item():.4f}")
결과
[Epoch 1] Batch 0, Loss: 2.2996
[Epoch 1] Batch 100, Loss: 0.5997
[Epoch 1] Batch 200, Loss: 0.4387
[Epoch 1] Batch 300, Loss: 0.4640
[Epoch 1] Batch 400, Loss: 0.1842
...
[Epoch 5] Batch 600, Loss: 0.0815
[Epoch 5] Batch 700, Loss: 0.0771
[Epoch 5] Batch 800, Loss: 0.0886
[Epoch 5] Batch 900, Loss: 0.1146
(2)CNN으로 MNIST를 학습시킨 경우
코드
# 0. Import
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
# 1. 데이터 변환 및 로딩
transform = transforms.ToTensor()
train_dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
# 2. CNN 모델 정의
class CNN(nn.Module):
def __init__(self): #layer 정의
super().__init__()
self.conv1 = nn.Conv2d( # (1) 2D Convolution: 공간 구조 유지하며 특징 추출
in_channels=1, # MNIST는 흑백이므로 채널 1개
out_channels=16, # 필터(커널) 16개 사용 → 특징맵 16개 생성
kernel_size=3, # 특징(3x3)필터
padding=1 # 가장자리 보존을 위해 zero-padding 1 적용
) # → 입력: [B, 1, 28, 28] → 출력: [B, 16, 28, 28]
self.pool = nn.MaxPool2d(2, 2) # (2) Max Pooling: 특징 유지하며 크기 줄이기
# kernel_size=2 : 2×2 영역을 한 블록으로 묶음
# tride=2 : 2칸씩 이동하며 풀링 → 겹치지 않음
# 연산 방식 : 2×2 안에서 최댓값 하나만 선택
# → [B, 16, 28, 28] → [B, 16, 14, 14]
self.fc1 = nn.Linear(16 * 14 * 14, 10) #최종 분류
# 16 * 14 * 14 : 이전 Conv + Pool 연산으로 얻은 이미지 특징의 총 개수
# 10 : 최종 클래스 수 (MNIST는 숫자 0~9 → 총 10개)
# 역할 : 16 * 16 * 14를 1차원으로 표서 10차원으로 변환(출력 10개 0~9로)
def forward(self, x): # 실행
x = self.pool(F.relu(self.conv1(x))) # 합성곱(특징 추출) + relu(특징 추출 도움) + pool로 단순화
# x.shape == [B, 16, 14, 14]
x = x.view(x.size(0), -1) # 평탄화(Flatten)
# size(0)은 B 즉 배치 크기 가져오고 -1은 나머지 숫자들 전부 곱하라는 뜻임
# [배치 크기, 특징 벡터 차원] → [64, 3136] 결국 이 형태로 됨
x = self.fc1(x) # Fully connected(완전 연결층으로 분류 마무리 0~9)
return x
model = CNN()
# 3. 손실 함수 & 옵티마이저
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
# 4. 학습 루프 (5 Epoch)
num_epochs = 5
for epoch in range(num_epochs):
for batch_idx, (images, labels) in enumerate(train_loader):
outputs = model(images)
loss = criterion(outputs, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if batch_idx % 100 == 0:
print(f"Batch {batch_idx}, Loss: {loss.item():.4f}")
결과
[Epoch 1] Batch 0, Loss: 2.2996
[Epoch 1] Batch 100, Loss: 0.5997
[Epoch 1] Batch 200, Loss: 0.4387
[Epoch 1] Batch 300, Loss: 0.4640
[Epoch 1] Batch 400, Loss: 0.1842
...
[Epoch 5] Batch 600, Loss: 0.1059
[Epoch 5] Batch 700, Loss: 0.1204
[Epoch 5] Batch 800, Loss: 0.0337
[Epoch 5] Batch 900, Loss: 0.0096
결론
MLP로 학습한 경우 최종 Loss값이 0.1 내외정도인 경우에 비해 CNN으로 학습한 경우 최종 Loss값이 0.01 내외로 CNN을 사용한 경우 Loss값이 더 적음을 알 수 있다. 즉 MNIST를 모델에 학습시킬 때 MLP보다 CNN인 경우 더 학습이 잘 되는 것을 확인할 수 있었다.