[Pytorch] 생초보의 파이토치 일기 - MNIST 손글씨 데이터 분류 99% 달성하기 1

 이번 포스팅에서는 파이토치를 거의 처음 써보는 제가 MNIST 손글씨 데이터를 분류하는 신경망을 파이토치로 구성하는 과정을 일기 형식으로 한 것입니다.

 

 그 동안 CVPR과 같은 메이져 논문을 읽으면서 느낀 것은 최근에 딥러닝 프레임워크 중 파이토치의 비중이 점점 늘어난 다는 점이다. 그에 반해 나는 딥러닝을 공부해본지 어언 1년 동안 오직 텐서플로우, 케라스만 잡고 있어 현재의 트렌드에 따라가기 많이 어려웠다. 

 

 일단 딥러닝을 시작하게 되면 항상 먼저 시작하는 과제 중에 하나가 MNIST 손글씨 데이터 셋이였기 때문에 이번에도 파이토치를 이용해서 진행해보기로 했다. 일단 최종 목표는 99% 달성이다. 그런데 최근 논문들의 코드들을 보면 tqdm과 같은 라이브러리를 추가적으로 사용해서 progress bar를 이쁘게 그리기도 했는 데 이 역시 연습해보기로 했다. 또한, 비록 간단한 모델일지라도 한 개의 py파일에 몰아서 넣는 것이 아니라 기능별로 정리하는 것도 연습해볼 예정이다. 

 

 일단 필요한 라이브러리부터 불러오도록 하자. 이 라이브러리들은 밥먹듯이 불러오는 라이브러리이기 때문에 외워두면 좋다고 한다. 

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision.transforms import transforms
from torchvision.datasets import MNIST

 먼저, "학습"을 하기 위해서는 무엇이 필요한 지 생각해보면 모델이 필요하고 모델에 주입할 데이터가 필요하다. 일단 데이터는 MNIST를 사용하기로 했기 때문에 파이토치의 공식 문서를 참조해보면 torchvision 라이브러리에 MNIST 데이터셋을 다운로드 받을 수 있는 코드가 존재한다.

torchvision.datasets.MNIST(root, train=True, transform=None, target_transform=None, download=False)

 root는 데이터셋을 다운받을 경로, train은 훈련 데이터를 받을 건지, 시험 데이터를 받을 건지를 결정한다. transform은 MNIST 데이터셋이 이미지이기 때문에 어떤 변환을 주어서 받을 것인지 결정하는 것이다. 나는 여기서 image를 tensor로 변환해주는 ToTensor와 이미지를 정규화하는 함수인 Normalize를 적용하였다. target_transform은 target 데이터, 즉 0~9까지 라벨링되어있는 데이터를 어떤 변화를 줄 것인지를 결정한다. download를 실제로 다운받을 것인지 결정한다. 따라서 이를 내 코드로 정리하면 아래와 같다. 

 

# MNIST 데이터의 transform 정의
mnist_transform = transforms.Compose([
    transforms.ToTensor(), # image -> tensor
    transforms.Normalize((0.1307, ), (0.3081, ))
])

# 데이터 로드하기
data_path = './MNIST_digit' # 다운로드할 폴더 선택
train_dataset = MNIST(data_path, transform=mnist_transform, train=True, download=True)
test_dataset = MNIST(data_path, transform=mnist_transform, train=False, download=True)

 나는 download를 True로 지정했기 때문에 이와 같이 MNIST 데이터가 따로 폴더가 생성되고 저장된 것을 볼 수 있다. 이때, processed 폴더는 transform을 적용한 데이터, raw는 적용하지 않은 데이터이다.

 

 자 이제 데이터를 불러왔으니 모델이 필요하다. 딥러닝에서 모델은 종류가 아주아주 많다. 대표적으로 DNN, CNN, RNN 등이 있다. 이번 포스팅에서는 DNN을 이용해서 구현을 해보도록 하겠다. 따라서, 입력 벡터의 사이즈는 28 $\times$ 2 = 784로 변환되어야한다. 그 다음으로 출력 벡터의 사이즈는 0~9까지의 digit을 예측하기 때문에 10으로 맞추어주면 된다. 따라서 내가 해주면 될 것은 hidden layer의 입력-출력의 크기를 정해주어야하는 데 처음 해보는 것이기 때문에 한 개의 hidden layer만 추가하고 입력은 256, 출력은 10로 정한다. 또한 layer의 사이사이에는 activation 함수를 걸어줘야하는 데 보통 Relu 함수를 많이 쓰고 다중 분류 문제에서는 마지막 activation 함수를 softmax로 걸기 때문에  input $\times$ hidden -> relu -> output -> softmax로 진행하도록 하겠다. 파이토치에서 신경망은 항상 클래스로 구현되며 이것은 대부분 형식이 정해져있기 때문에 외워두는 것이 편할 거 같다. 실제로 보면 매우 단순한 신경망이다.

# 신경망 클래스 정의
class Network(nn.Module) :
    def __init__(self):
        super(Network, self).__init__()

        self.fc1 = nn.Linear(784, 256)
        self.fc2 = nn.Linear(256, 10)

    def forward(self, inputs):
        x = inputs.view(-1, 784)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

 그 다음으로는 데이터를 신경망으로 넘겨주는 Data loader라는 것을 정의해주어야한다. 간단히 말해서 데이터와 신경망 사이의 도로라고 생각하면 될 거 같다. 어떤 데이터를 넘겨주는 지와 배치 사이즈만 정해주면 되며, 만약 shuffle를 True로 하게 되면 한번 배치 사이즈만큼 데이터를 로드할 때 한번 섞기 때문에 동일한 데이터가 여러번 뽑힐 수도 있다.

# 데이터 로더 정의
train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=True)

 이제부터는 공식이다. 정의한 모델을 인스턴스 -> 최적화 알고리즘 정의 -> 학습 시작. 만약 GPU가 있다면 한줄의 코드만 추가해주면 된다.

# 모델 인스턴스
model = Network()
print("모델 정의 완료!")

# GPU로 넘기기
model = model.to(device)

# 최적화 알고리즘 정의
optimizer = optim.SGD(params=model.parameters(), lr=lr)

 이제부터는 학습하는 코드이다. 미리 정의해둔 에폭 횟수만큼 학습을 진행하게 된다. 이때, 에폭이란 정해진 배치 사이즈로 전체 데이터 셋을 한번 싹 훑은 것을 의미한다. 나는 여기서 배치 사이즈를 128개로 정의했기 때문에 내가 정의한 Data Loader는 60000 / 128 = 468.75로 468번 데이터를 로드한다. 따라서 여기서 1에폭=데이터 로더가 468번 데이터를 로드함을 의미한다.

 

 데이터를 한번 로드하면 제일 중요한 것은 optimizer의 그래디언트를 0으로 초기화를 한번 해주어야한다는 점이다. 만약 해주지 않으면 버퍼가 걸려서 잘 안된다고 한다. 이유는 잘모르겠다... 하지만 꼭 초기화 해주라고 하니 주의하자. 그리고 forward propagation -> 손실 계산 -> backward propagation -> 가중치 업데이트 순으로 진행된다. 1에폭을 학습할 때마다 시험 데이터를 이용해서 손실과 정확도를 계산할 수 있다. 이때, 반드시 with 블록을 함께 추가하여 파이토치가 시험 데이터의 손실을 계산하는 데 그래디언트를 계산하지 않도록 주의하자. 따라서 forward propagation -> 손실 계산만 하면 시험 데이터는 종료된다.

# train
train_loss_list, test_loss_list = [], []
for epoch in range(epochs) :
    model.train() # 모델을 train 모드로 변환
    train_loss = 0
    for batch_idx, (data, label) in enumerate(train_loader) :
        print(batch_idx)
        data, label = data.to(device), label.to(device)
        optimizer.zero_grad() # 그래디언트 초기화
        output = model(data) # forward propagation
        loss = F.nll_loss(output, label) # 손실 계산
        loss.backward() # backward propagation
        optimizer.step() # 가중치 업데이트
        if batch_idx % 50 == 0 :
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()))
        train_loss += loss.item()
    train_loss_list.append(train_loss/batch_size)

    model.eval() # 모델을 eval 모드로 변환
    test_loss = 0
    correct = 0
    with torch.no_grad() :
        for data, label in test_loader :
            data, label = data.to(device), label.to(device)
            output = model(data)
            test_loss += F.nll_loss(output, label, reduction='sum').item()
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(label.view_as(pred)).sum().item()

        test_loss /= len(test_loader.dataset)
        test_loss_list.append(test_loss)

        print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format
              (test_loss, correct, len(test_loader.dataset),
               100. * correct / len(test_loader.dataset)))

나는 총 20에폭을 돌려서 약 96%의 성능을 얻었다. 나름 나쁘지 않은 성능이지만 내 목표는 99%이상이기 때문에 다른 방법을 찾도록 해보자.

 

이번에는 기존의 DNN 모델을 조금 더 깊게 쌓아서 구현하였다. 

# 신경망 클래스 정의
class Network(nn.Module) :
    def __init__(self):
        super(Network, self).__init__()

        self.fc1 = nn.Linear(784, 256)
        self.fc2 = nn.Linear(256, 128)
        self.fc3 = nn.Linear(128, 64)
        self.fc4 = nn.Linear(64, 10)

    def forward(self, inputs):
        x = inputs.view(-1, 784)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = F.relu(self.fc3(x))
        x = self.fc4(x)
        return F.log_softmax(x, dim=1)

2개의 hidden layer를 추가한 모습니다. 그 결과 97%로 얕은 모델을 사용했을 때 보다 1% 향상된 성능을 얻을 수 있었다. 

 

다음은 CNN 모델을 만들어서 구현해보자. 나머지는 동일하고 신경망만 바꾸면 된다.

# 신경망 클래스 정의
class Network(nn.Module) :
    def __init__(self):
        super(Network, self).__init__()

        # input shape = (?, 28, 28, 1)
        #   Conv     -> (?, 28, 28, 32)
        #   Pool     -> (?, 14, 14, 32)
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
        self.maxpool1 = nn.MaxPool2d(kernel_size=2, stride=2)

        # input shape = (?, 14, 14, 32)
        #   Conv     -> (?, 14, 14, 64)
        #   Pool     -> (?, 7, 7, 64)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
        self.maxpool2 = nn.MaxPool2d(kernel_size=2, stride=2)

        self.fc1 = nn.Linear(7*7*64, 256)
        self.fc2 = nn.Linear(256, 128)
        self.fc3 = nn.Linear(128, 64)
        self.fc4 = nn.Linear(64, 10)
    def forward(self, inputs):
        x = F.relu(self.conv1(inputs))
        x = self.maxpool1(x)
        x = F.relu(self.conv2(x))
        x = self.maxpool2(x)

        x = x.view(x.size(0), -1) # Flatten
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = F.relu(self.fc3(x))
        x = self.fc4(x)

        return F.log_softmax(x, dim=1)

 아까보다 살짝 복잡하지만 CNN 계층만 2개 추가한 것이다. 그 결과 99%로 최종적으로 내가 원하는 성능이 나왔다!! 다음 포스팅에서는 코드를 최종적으로 깔끔하게 정리하는 연습을 해보겠다.