PyTorch를 이용한 CIFAR-10 데이터 세트 학습 코드 및 결과 총 정리 (ResNet 활용하여 95% 이상 정확도 얻기)

※ 본 게시글은 필자의 개인 경험을 토대로 작성된 것으로, 잘못된 정보를 포함하고 있을 수 있습니다. ※

 

  NVIDIA Tesla V100 혹은 NVIDIA TITAN RTX과 같은 GPU 하나만 있어도 2시간 정도면 CIFAR-10 데이터 세트를 학습하여 Top-1 정확도(accuracy)로 95% 이상을 뽑아낼 수 있다. 이번 포스팅에서는 간단히 PyTorch Hub에서 기본적으로 제공하는, 흔히 알려진 CIFAR-10 데이터 세트에 대하여 CNN 모델(ResNet-18)을 학습을 진행하는 방법에 대해서 알아보겠다. 구체적으로 ResNet-18 모델을 사용하여 간단히 학습을 진행해 볼 수 있다.

 

  가장 먼저, Jupyter Notebook 상에서 다음과 같이 사용할 GPU의 번호를 설정할 수 있다.

 

import os
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "6"

 

  CIFAR-10의 경우 데이터 세트의 크기가 작기 때문에, PyTorch에서 기본적으로 제공하고 있는 것을 확인할 수 있다.

 

import torchvision.datasets as datasets
import torchvision.transforms as transforms
import torchvision

batch_size = 256

augment_train = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
])

transform_train = transforms.Compose([
    transforms.ToTensor(),
])

train_dataset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_train)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=20)

transform_test = transforms.Compose([
    transforms.ToTensor(),
])

test_dataset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=20)

 

  필자는 다음과 같은 방식으로 학습을 진행했다. 참고로 코드를 잘 살펴보면, 일단은 PyTorch를 텐서(tensor) 형태로 불러와 배치 단위로 확인한 뒤에, 이를 다시 PIL 이미지로 바꾸어서 내가 원하는 데이터 증진(augmentation) 기법을 적용하는 코드 형식으로 작성된 것을 확인할 수 있다.

 

def train(epoch):
    net.train()
    train_loss = 0
    correct = 0
    total = 0
    for batch_idx, (images, labels) in enumerate(train_loader):
        images = images.cuda()
        targets = labels.cuda()

        augmented_images = torch.FloatTensor([])
        for image in images:
            pil_image = to_pil_image(image)
            augmented_images = torch.cat([augmented_images, augment_train(pil_image).unsqueeze(0)], dim=0)
        augmented_images = augmented_images.cuda()
        
        optimizer.zero_grad()

        benign_outputs = net(augmented_images)
        loss = criterion(benign_outputs, targets)
        loss.backward()

        optimizer.step()
        train_loss += loss.item()
        _, predicted = benign_outputs.max(1)

        total += targets.size(0)
        correct += predicted.eq(targets).sum().item()

    print(f'[{epoch}] Total benign train accuarcy:', 100. * correct / total)
    return 100. * correct / total


def test(epoch):
    net.eval()
    benign_loss = 0
    benign_correct = 0
    total = 0
    with torch.no_grad():
        for batch_idx, (inputs, targets) in enumerate(test_loader):
            inputs, targets = inputs.cuda(), targets.cuda()
            total += targets.size(0)

            outputs = net(inputs)
            loss = criterion(outputs, targets)
            benign_loss += loss.item()

            _, predicted = outputs.max(1)
            benign_correct += predicted.eq(targets).sum().item()

    print(f'[{epoch}] Total benign test accuarcy:', 100. * benign_correct / total)
    return 100. * benign_correct / total

 

  또한, 학습(training) 코드와 테스트(test) 코드가 각각 함수로 나누어져 작성되어 있기 때문에, 체계적으로 로깅(logging)을 진행하면서 학습이 되는 과정을 기록할 수 있을 것이다. 필자는 다음과 같이 logger 객체를 사용했다.

 

from torchvision.transforms.functional import to_pil_image
import logging

logger = logging.getLogger()
logger.setLevel(logging.INFO)

log_path = "training.log"
handler = logging.FileHandler(log_path, 'a')
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
handler.setFormatter(formatter)

logger.addHandler(handler)

 

  다음과 같이 학습을 해볼 수 있다. 필자는 0.1부터 학습률(learning rate)이 시작되어, 코사인 학습률(cosine annealing learning rate) 기법을 적용하여 서서히 학습률이 감소하도록 하였다. 이는 일반적으로 많이 사용되는 하이퍼 파라미터 세팅이기도 하다. 결과적으로, 다음과 같은 코드를 작성할 수 있다.

 

import torch.optim as optim

criterion = nn.CrossEntropyLoss()
learning_rate = 0.1
net = ResNet18(channel=3, num_classes=10, im_size=(32, 32, 3))
net = net.cuda()

optimizer = optim.SGD(net.parameters(), lr=learning_rate, momentum=0.9, weight_decay=0.0005)
num_epochs = 200
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_epochs, eta_min=0.0001)
best_accuracy = 0

start_time = time.time()
for epoch in range(0, num_epochs):
    train_accuracy = train(epoch)
    test_accuracy = test(epoch)
    best_accuracy = max(best_accuracy, test_accuracy)
    scheduler.step()

    message = f"[{epoch}] train_accuracy: {train_accuracy:.2f}%, test_accuracy: {test_accuracy:.2f}%, elapsed time: {time.time() - start_time:.4f} seconds."
    print(message)
    logger.info(message)

 

  로그 파일 "training.log" 파일의 내용으로는 바로 다음과 같이 기록되는 것을 확인할 수 있었다. (일부 생략된 결과) 그렇다면 얼마나 많은 시간이 소요되는가? 최종적으로 200번의 epoch을 위해 약 2시간 정도가 소요되는 것을 확인할 수 있으며, 200번의 epoch을 반복했을 때 최종적인 정확도(accuracy)는 95%가 되는 것을 확인할 수 있다.

 

INFO - [0] train_accuracy: 24.08%, test_accuracy: 27.68%, elapsed time: 44.2970 seconds.
INFO - [1] train_accuracy: 40.68%, test_accuracy: 22.61%, elapsed time: 87.2885 seconds.
INFO - [2] train_accuracy: 51.72%, test_accuracy: 48.65%, elapsed time: 131.4002 seconds.
INFO - [3] train_accuracy: 60.89%, test_accuracy: 60.80%, elapsed time: 175.8557 seconds.
INFO - [4] train_accuracy: 67.16%, test_accuracy: 47.24%, elapsed time: 218.8372 seconds.
INFO - [5] train_accuracy: 72.14%, test_accuracy: 68.02%, elapsed time: 262.8409 seconds.
INFO - [6] train_accuracy: 75.94%, test_accuracy: 66.05%, elapsed time: 305.8487 seconds.
INFO - [7] train_accuracy: 78.86%, test_accuracy: 71.59%, elapsed time: 349.2591 seconds.
INFO - [8] train_accuracy: 80.85%, test_accuracy: 76.81%, elapsed time: 392.8524 seconds.
INFO - [9] train_accuracy: 82.68%, test_accuracy: 78.71%, elapsed time: 438.7123 seconds.
INFO - [19] train_accuracy: 89.20%, test_accuracy: 81.09%, elapsed time: 877.2395 seconds.
INFO - [29] train_accuracy: 91.33%, test_accuracy: 84.95%, elapsed time: 1313.6072 seconds.
INFO - [39] train_accuracy: 92.47%, test_accuracy: 85.61%, elapsed time: 1752.1859 seconds.
INFO - [49] train_accuracy: 93.48%, test_accuracy: 82.64%, elapsed time: 2192.1826 seconds.
INFO - [59] train_accuracy: 93.93%, test_accuracy: 87.73%, elapsed time: 2632.6115 seconds.
INFO - [69] train_accuracy: 94.64%, test_accuracy: 88.19%, elapsed time: 3070.4473 seconds.
INFO - [79] train_accuracy: 95.25%, test_accuracy: 88.18%, elapsed time: 3507.3784 seconds.
INFO - [89] train_accuracy: 96.17%, test_accuracy: 90.65%, elapsed time: 3947.9512 seconds.
INFO - [99] train_accuracy: 96.45%, test_accuracy: 89.65%, elapsed time: 4389.4442 seconds.
INFO - [109] train_accuracy: 97.39%, test_accuracy: 88.95%, elapsed time: 4827.5396 seconds.
INFO - [119] train_accuracy: 97.81%, test_accuracy: 91.03%, elapsed time: 5269.8396 seconds.
INFO - [129] train_accuracy: 98.52%, test_accuracy: 92.44%, elapsed time: 5711.2712 seconds.
INFO - [139] train_accuracy: 99.34%, test_accuracy: 93.05%, elapsed time: 6148.9241 seconds.
INFO - [149] train_accuracy: 99.70%, test_accuracy: 94.14%, elapsed time: 6592.9668 seconds.
INFO - [159] train_accuracy: 99.97%, test_accuracy: 95.04%, elapsed time: 7031.4187 seconds.
INFO - [169] train_accuracy: 100.00%, test_accuracy: 95.08%, elapsed time: 7475.8688 seconds.
INFO - [179] train_accuracy: 100.00%, test_accuracy: 95.19%, elapsed time: 7916.4395 seconds.
INFO - [189] train_accuracy: 100.00%, test_accuracy: 95.24%, elapsed time: 8359.8897 seconds.
INFO - [199] train_accuracy: 100.00%, test_accuracy: 95.19%, elapsed time: 8797.8429 seconds.

 

  또한, 10번 정도의 epoch만을 반복했음에도 테스트 정확도(test accuracy)가 순식간에 80%가 되는 것을 확인할 수 있다.