PyTorch에서 특정 Dataset을 열어 이미지 출력하기

  실험을 하면서 자주 쓰는 코드인데, 따로 정리를 해놓지 않아서 매 번 입력을 하고 있다. 그래서 정리하려고 한다. 일단 Dataset 객체를 불러올 때는 데이터를 전처리하는 부분이 들어간다. PyTorch의 경우 ToTenser() 함수를 불러오면, 이미지가 자동으로 [0, 1]의 값으로 변경된다. 예를 들어 CIFAR-10 학습용 데이터셋을 불러오는 코드는 다음과 같다.

 

import torch
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np

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=128, shuffle=True, num_workers=4)

 

  이렇게 불러온 이미지는 실제로 Tensor 객체로 존재하며, 각 원소의 값이 0부터 1 사이의 값이다. 그렇기 때문에 이를 화면에 출력하고자 한다면, 이 값을 다시 0부터 255 사이의 값으로 늘려야 하는 건지 궁금할 수 있다. 다행히도 파이썬의 matplotlib는 기본적으로 0부터 1사이의 값이라고 해도 알아서 인식하여 정상적인 이미지로 출력해준다. 하지만 별도로 OpenCV 등에서 활용하고자 한다면, 추가적인 전처리가 필요할 수 있다.

 

  또한 기본적으로 PyTorch는 이미지 데이터셋을 [Batch Size, Channel, Width, Height] 순서대로 저장하기 때문에, 이를 matplotlib로 출력하기 위해서는 각 이미지를 [Width, Height, Channel] 형태로 변경해 줄 필요가 있다. 그것은 numpy 라이브러리의 transpose() 함수를 이용하여 해결할 수 있다.

 

def custom_imshow(img):
    img = img.numpy()
    plt.imshow(np.transpose(img, (1, 2, 0)))
    plt.show()

 

  그럼 이제 이렇게 정의된 imshow() 함수를 이용하여 데이터셋의 이미지를 배치당 하나씩 출력해보자.

 

def process():
    for batch_idx, (inputs, targets) in enumerate(train_loader):
        custom_imshow(inputs[0])

process()

 

  실행 결과는 다음과 같다.