PyTorch 기초 - 전체적인 학습 코드 정리

이전 포스트까지는 딥러닝에서 학습을 수행하기 위한 각각의 클래스, 함수 등을 살펴보았습니다. 이 포스트에서는 복습하는 느낌으로 이를 종합해서, 간단한 MNIST 분류기를 만들어보겠습니다.

MNIST 데이터셋 살펴보기

MNIST 데이터셋는 0에서 9까지의 직접 쓰여진 이미지와 해당되는 숫자로 이루어져 있습니다. 일부 데이터들을 통해서 어떤 느낌으로 이루어져 있는지 살펴보겠습니다.

import matplotlib.pyplot as plt

transform = transforms.ToTensor()
train_data = datasets.MNIST(
    root="./data", train=True, download=True, transform=transform
)

plt.rcParams["figure.figsize"] = (10, 12)
for i in range(100):
    plt.subplot(10, 10, i + 1)
    plt.imshow(train_data[i][0][0], cmap="gray")
    plt.title(train_data[i][1])
    plt.axis("off")
plt.show()

눈으로 어떤 이미지를 보았을 때 이게 어떤 숫자인지 유추하기는 쉬울 것입니다. 그런데 배열로 이루어져 있는 컴퓨터의 시각으로 보았을 때에는 그렇지 않을 수도 있습니다. 이 문제를 해결하기 위해 CNN 구조의 모델을 통과하는 방식을 사용해서, 어떤 이미지가 주어졌을 때 해당하는 숫자에 해당하는지 유추하는 분류기를 만들 수 있을 것입니다. 이 데이터셋은 듀토리얼용으로 아주 유명한 예시라, 코드에서도 명시했다 싶이 torchvision.datasets.MNIST로 바로 가져올 수 있습니다. 이 클래스를 활용해 데이터를 가져올 때 별도의 transform 명시가 없으면 PIL로 가져오기 때문에 텐서로 바꾸어서 가져오는 것이 좋습니다.

MNIST 분류기 만들기

그러면 본격적으로 CNN 모델을 활용해서 MNIST 분류기를 만들고, 이를 학습하는 코드를 구현해보겠습니다.

# 1
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

# 2
def SimpleCNNLayer(in_features, out_features):
    return nn.Sequential(
        nn.Conv2d(in_features, out_features, 3, 1, 1),
        nn.BatchNorm2d(out_features),
        nn.ReLU(),
        nn.MaxPool2d(2),
    )


def SimpleCNNClassifier():
    return nn.Sequential(
        SimpleCNNLayer(1, 32),
        SimpleCNNLayer(32, 32),
        nn.Flatten(),
        nn.Linear(1568, 64),
        nn.Linear(64, 10),
    )


# 3
transform = transforms.ToTensor()

# 4
train_data = datasets.MNIST(
    root="./data", train=True, download=True, transform=transform
)
test_data = datasets.MNIST(
    root="./data", train=False, download=True, transform=transform
)
train_loader = DataLoader(train_data, batch_size=32, shuffle=True)
test_loader = DataLoader(test_data, batch_size=32, shuffle=False)

# 5
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 6
model = SimpleCNNClassifier().to(device)

# 7
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

for epoch in range(10):
    # 8
    model.train()
    for images, labels in train_loader:
        # 9
        images = images.to(device)
        labels = labels.to(device)

        # 10
        outputs = model(images)
        loss = criterion(outputs, labels)

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

    # 12
    model.eval()
    test_loss = 0.0
    correct = 0

    # 13
    with torch.no_grad():
        for images, labels in test_loader:
            images = images.to(device)
            labels = labels.to(device)

            # 14
            outputs = model(images)
            predicted = torch.max(outputs, 1)[1]
            loss = criterion(outputs, labels)

            # 15
            test_loss += loss.item()
            correct += (labels == predicted).sum()
    # 16
    print(
        f"epoch {epoch+1} - test loss: {test_loss / len(test_loader):.4f}, accuracy: {correct / len(test_data):.4f}"
    )

학습을 위한 초기화 과정과 학습할 때의 전체적인 구조는 앞으로도 비슷할 것이기 때문에, 전체적인 흐름을 이해하는 것이 중요하겠습니다. 그러면 순서대로 코드를 살펴보겠습니다.

1. 사용할 모듈을 불러옵니다.

2. 사용할 CNN 모델을 nn.Sequential을 활용해서 구현합니다.

3. MNIST 데이터셋을 가져오고, 적용할 transform을 명시합니다. 여기서는 별도의 augmentation을 적용하지는 않고, PIL 인스턴스에서 텐서로 만들기 위해 ToTensor()를 사용하도록 합니다.

4. Dataset과 DataLoader를 정의합니다. 여기서는 MNIST 데이터셋을 불러오기 위해 torchvision.datasets.MNIST를 사용하였습니다. DataLoader는 배치 사이즈, 셔플 등을 명시해서 만들어줍니다.

5. CPU나 GPU를 사용할 것인지 device를 정의합니다. 여기서는 GPU CUDA를 사용할 수 있을 때 cuda를 사용하도록 합니다.

6. 모델을 인스턴스로 정의합니다. 여기서 .to(device)가 붙었는데, 이는 모델(가중치)을 사용할 장치로 올려놓는 역할을 합니다. 특히 GPU를 사용할 때에는 꼭 사용하셔야 합니다.

7. 손실 함수와 optimizer를 정의합니다.

8. 이제 학습 코드로 들어왔습니다. 데이터셋에 대해 반복할 수를 에포크로 명시하고, 학습을 수행할 때 모델을 .train()으로 학습 상태로 명시해줍니다. 참고로 학습와 평가 상태에서는 드롭아웃과 배치 정규화 레이어에서의 동작에 차이가 있습니다.

9. DataLoader로부터 이미지와 해당하는 숫자를 묶은 텐서를 images, labels로 가져와서, 사용할 장치로 데이터를 보냅니다.

10. 입력 이미지에 대해 모델을 통과시킵니다. 그리고, 나온 출력 텐서와 타겟 텐서 간 손실 함수를 통해 손실을 계산합니다.

11. 역전파 알고리즘을 통해 기울기를 계산하고, optimizer를 활용해서 모델의 가중치를 조정합니다.

12. 모델을 .eval()으로 평가 상태로 명시해줍니다. 그리고 평가한 성능을 기록하기 위한 변수들을 선언해줍니다.

13. torch.no_grad()는 모델을 통과하는 데 기울기를 트래킹하지 않는다는 것을 명시해줍니다. 이를 사용했을 때 트래킹에 사용되는 메모리 사용량이나 연산을 줄일 수 있습니다.

14. 테스트 이미지에 대해 모델을 통과시켜서 손실을 계산합니다. 그리고 나온 10개의 숫자에 대해 추정한 값 중 가장 높은 숫자로 판정하도록 합니다.

15. 각 에포크마다 손실이 얼마나 감소하는지, 얼마나 맞추었는지 확인하기 위해 기록합니다.

16. 기록한 손실과 정확도에 대해 출력합니다.

이렇게 데이터 로딩, 손실 함수와 optimizer 선정, 학습 과정에서 순방향 통과와 역방향 전파, 평가 과정을 통한 성능 확인 과정을 코드를 통해 간단하게 살펴보았습니다. 이 코드를 실행했을 때 나오는 예시 결과입니다.

epoch 1 - test loss: 0.0581, accuracy: 0.9811
epoch 2 - test loss: 0.0476, accuracy: 0.9853
epoch 3 - test loss: 0.0435, accuracy: 0.9872
epoch 4 - test loss: 0.0311, accuracy: 0.9904
epoch 5 - test loss: 0.0320, accuracy: 0.9893
epoch 6 - test loss: 0.0330, accuracy: 0.9903
epoch 7 - test loss: 0.0352, accuracy: 0.9885
epoch 8 - test loss: 0.0355, accuracy: 0.9898
epoch 9 - test loss: 0.0422, accuracy: 0.9883
epoch 10 - test loss: 0.0431, accuracy: 0.9885

4 에포크까지 테스트 손실이 감소하고 정확도가 증가하다가, 그 이후로는 조금씩 떨어지는 것을 볼 수 있습니다. 이를 참고해서 좋은 분류기를 만들도록 조정을 해주면 됩니다.