CNN으로 MNIST 분류하기

출처 : https://coding-yesung.tistory.com/223

 

• Import 하기

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

 

• cuda 데꼬오기

device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(device)

 

• 데이터 불러오기

train_data = datasets.MNIST(
    root = 'data',   # 데이터셋이 저장될 디렉토리 경로
    train=True,      # 훈련 데이터를 가져옴
    transform=transforms.ToTensor(),  # 이미지 데이터를 텐서로 변환하는 변환
    download=True   # 데이터셋이 로컬에 없을 경우 인터넷에서 데이터를 자동으로 다운로드
)
test_data = datasets.MNIST(
    root = 'data',
    train=False,    # test용으로 설정하여 로드
    transform=transforms.ToTensor(),  # 이미지 데이터를 텐서로 변환하는 변환
    download=True
)


print(train_data)
print(test_data)

 

 

• Loader

loader = DataLoader(
    dataset = train_data,   # 훈련 데이터셋을 배치단위로 로드
    batch_size=64,  # 배치 크기 64
    shuffle=True  # 데이터를 섞어줌
)

 

 

• 데이터 확인

imgs, labels = next(iter(loader))  # 데이터로더에서 배치를 가져옴

fig, axes = plt.subplots(8, 8, figsize=(16, 16))  # 16*16픽셀의 이미지데이터가 8행 8열로 만드는 subplot

for ax, img, label in zip(axes.flatten(), imgs, labels):
  ax.imshow(img.reshape((28,28)), cmap='gray')  #  이미지를 28x28 크기로 변형, 흑백으로 표시
  ax.set_title(label.item())  # lable값을 제목으로 설정
  ax.axis('off')  # 축 숨김

 

 

• 모델 만들기

model = nn.Sequential(
    nn.Conv2d(1, 32, kernel_size=3, padding='same'),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2),

    nn.Conv2d(32, 64, kernel_size=3, padding='same'),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2),

    # Flatten
    nn.Flatten(),
    nn.Linear(7 * 7 * 64, 10)
).to(device)  # gpu로 보냄

print(model)

 

• Train 진행

optimizer = optim.Adam(model.parameters(), lr=0.001)

epochs = 10

for epoch in range(epochs+1):
  sum_losses = 0
  sum_accs = 0

  for x_batch, y_batch in loader:
    x_batch = x_batch.to(device)  # gpu 연산을 위해 보냄
    y_batch = y_batch.to(device)

    y_pred = model(x_batch)

    loss = nn.CrossEntropyLoss()(y_pred, y_batch)  # 3개 이상 클래스로 분류 -> CrossEntropyLoss

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    # 배치 단위 loss 저장
    sum_losse = sum_losses + loss.item()
    # 배치 단위 정확도 저장
    y_prob = nn.Softmax(1)(y_pred)
    y_pred_index = torch.argmax(y_prob, axis=1)
    acc = (y_batch == y_pred_index).float().sum() / len(y_batch) * 100

    sum_accs = sum_accs + acc

  avg_loss = sum_losses / len(loader)
  avg_acc = sum_accs / len(loader)

  print(f'Epoch: {epoch+1:4d}/{epochs}  Loss: {avg_loss:6f} Accuracy: {avg_acc:2f}%')

 

• 테스트 데이터로더

test_loader = DataLoader(
    dataset = test_data,
    batch_size=64,
    shuffle=True
)

 

 

• 테스트 데이터 확인

imgs, labels = next(iter(test_loader))

fig, axes = plt.subplots(8, 8, figsize=(16, 16))  # 16*16픽셀의 이미지데이터가 8행 8열로 만드는 subplot

for ax, img, label in zip(axes.flatten(), imgs, labels):
  ax.imshow(img.reshape((28,28)), cmap='gray')
  ax.set_title(label.item())
  ax.axis('off')

 

• Evaluate

model.eval()  # 모델을 테스트 모드로 전환(학습시키려고 모델을 사용하는게 아니라 평가모드로 전환)

sum_accs = 0


for x_batch, y_batch in test_loader:
  x_batch = x_batch.to(device)  # gpu 연산을 위해 보냄
  y_batch = y_batch.to(device)

  y_pred = model(x_batch)

  y_prob = nn.Softmax(1)(y_pred)
  y_pred_index = torch.argmax(y_prob, axis=1)
  acc = (y_batch == y_pred_index).float().sum() / len(y_batch) * 100

  sum_accs = sum_accs + acc

avg_loss = sum_losses / len(test_loader)
avg_acc = sum_accs / len(test_loader)

print(f'테스트 정확도는 {avg_acc:.2f}% 입니다')