pytorch로 MNIST 분류하기

데이터 transform 정의

• 데이터를 불러오기 전에 정규화에 대한 기준을 설정해준다.
  평균: 0.5, 표준편차: 1.0

import torchvision.transforms as transforms

# Normalize data with mean=0.5, std=1.0
mnist_transform = transforms.Compose([
    transforms.ToTensor(), 
    transforms.Normalize((0.5,), (1.0,))
])

 

 

데이터 불러오기

• torchvision에 내장되어 있는 MNIST 데이터 불러오기

from torchvision.datasets import MNIST

# download path 정의
download_root = './MNIST_DATASET'

train_dataset = MNIST(download_root, transform=mnist_transform, train=True, download=True)
valid_dataset = MNIST(download_root, transform=mnist_transform, train=False, download=True)
test_dataset = MNIST(download_root, transform=mnist_transform, train=False, download=True)

데이터 로딩

• DataLoader 활용

from torch.utils.data import DataLoader

# option 값 정의
batch_size = 1

train_loader = DataLoader(dataset=train_dataset, 
                         batch_size=batch_size,
                         shuffle=True)

valid_loader = DataLoader(dataset=test_dataset, 
                         batch_size=batch_size,
                         shuffle=True)

test_loader = DataLoader(dataset=test_dataset, 
                         batch_size=batch_size,
                         shuffle=True)

classes 정의

classes = ('0', '1', '2', '3',
           '4', '5', '6', '7', '8', '9')

이미지 확인해보기

import matplotlib.pyplot as plt
import numpy as np

# 이미지를 보여주기 위한 함수

def imshow(img):
    img = img / 2 + 0.5     # unnormalize
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()

# 학습용 이미지를 무작위로 가져오기
dataiter = iter(train_loader)
images, labels = dataiter.next()

print(len(images), len(labels))

# 이미지 보여주기
imshow(torchvision.utils.make_grid(images))
# 정답(label) 출력
print(' '.join('%5s' % classes[labels[j]] for j in range(1)))

합성곱 신경망 정의하기

import torch.nn as nn
import torch.nn.functional as F


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 3, 2)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(3, 3, 3)
        self.fc1 = nn.Linear(1 * 3 * 5 * 5, 30)
        self.fc2 = nn.Linear(30, 15)
        self.fc3 = nn.Linear(15, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 1 * 3 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


net = Net()

손실함수와 Optimizer 정의하기

• 교차 엔트로피 손실(Cross-Entropy loss)과 모멘텀(momentum) 값을 갖는 SGD를 사용합니다.

  import torch.optim as optim
  

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

# 신경망 학습하기

-   이제 재미있는 부분이 시작됩니다. 단순히 데이터를 반복해서 신경망에 입력으로 제공하고, 최적화(Optimize)만 하면 됩니다.  

a = 0
for epoch in range(1): # 데이터셋을 수차례 반복합니다.

running_loss = 0.0
for i, data in enumerate(train_loader, 0):
    # [inputs, labels]의 목록인 data로부터 입력을 받은 후;
    inputs, labels = data

    # 변화도(Gradient) 매개변수를 0으로 만들고
    optimizer.zero_grad()

    # 순전파 + 역전파 + 최적화를 한 후
    outputs = net(inputs)
    loss = criterion(outputs, labels)
    loss.backward()
    optimizer.step()

    # 통계를 출력합니다.
    running_loss += loss.item()
    if i % 1000 == 999:    # print every 2000 mini-batches
        print('[%d, %5d] loss: %.3f' %
              (epoch + 1, i + 1, running_loss / 1000))
        running_loss = 0.0
    a = a + 1

    if a == 30000:
        break

print('Finished Training')

# 학습된 모델 저장하기

-   학습한 모델을 저장해보겠습니다.  

PATH = './mnist_net.pth'
torch.save(net.state_dict(), PATH)

# 시험용 데이터로 신경망 검사하기

dataiter = iter(test_loader)
images, labels = dataiter.next()

이미지를 출력합니다.

imshow(torchvision.utils.make_grid(images))
print('GroundTruth: ', ' '.join('%5s' % classes[labels[j]] for j in range(1)))

```

이렇게 하게 되면 Pytorch로 MNIST를 분석할 수 있습니다.