[PyTorch] PyTorch Tutorial 공부하기 #2 (Dataset, DataLoader)

본 글은 파이토치 한국어 튜토리얼을 공부하며 작성했음을 밝힙니다.

더 좋은 가독성(readability)과 모듈성(modularity)을 위해 데이터셋 코드를 모델 학습 코드로부터 분리하는것이 이상적이다.

이에, PyTorch는 torch.utils.data.DataLoader와 torch.utils.data.Dataset의 두 가지 데이터 기본 요소를 제공하여

1. 미리 준비해둔(pre-loaded) 데이터셋 (like MNIST)

2. 가지고 있는 데이터

를 사용할 수 있도록 한다.

 

• Dataset : 샘플과 정답(label)을 저장
• DataLoader : Dataset을 각 샘플에 쉽게 접근할 수 있도록 순회 가능한 객체로 감싸놓음

 1. 데이터셋 불러오기

TorchVision에서 MNIST 데이터셋을 불러오는 예제를 작성해보자.

MNIST 데이터는 손글씨 이미지 데이터셋으로 60,000개의 학습 예제와 10,000개의 테스트 예제로 이루어져 있다.

각 예제는 흑백(grayscale)의 28X28 이미지와 10개의 부류(class) 중 하나인 정답(label)로 구성된다.

 

• root : 학습/테스트 데이터가 저장되는 경로
• train : 학습용(True) 또는 테스트용(False) 데이터셋 여부를 결정함
• download : True로 설정하면 root에 데이터가 없는 경우 인터넷에서 다운로드함
• transform, target_transform : 특징(feature)과 정답(label)의 변형(transform)을 지정함

import torch
from torch.utils.data import Dataset
from torchvision import datasets
from torchvision.transforms import ToTensor
import matplotlib.pyplot as plt

training_data = datasets.MNIST(
    root = '../datasets',
    train = True,
    download = True,
    transform = ToTensor()
)

test_data = datasets.MNIST(
    root = '../datasets',
    train = False,
    download = True,
    transform = ToTensor()
)

위의 코드에서 ToTensor()는 PIL Image나 Numpy의 ndarray를 FloatTensor 형태로 변환하고, 이미지의 각 픽셀값을 [0, 1] 범위로 정규화하는 작업을 한다.

 

 2. DataLoader로 학습용 데이터 준비하기

Dataset은 데이터셋의 특징(feature)을 가져오고 하나의 샘플에 정답(label)을 지정하는 일을 한 번에 한다.

모델을 학습할 때, 일반적으로 샘플들을 미니배치(minibatch)로 전달하고, 매 epoch마다 데이터를 다시 섞어서 과적합(overfit)을 막고, multiprocessing을 사용해 데이터 검색 속도를 높인다.

 

DataLoader는 간단한 API로 이러한 복잡한 과정들을 추상화한 순회 가능한 객체다.

from torch.utils.data import DataLoader

train_dataloader = DataLoader(training_data, batch_size = 64, shuffle = True)
test_dataloader = DataLoader(test_data, batch_size = 64, shuffle = True)

 

 3. DataLoader로 순회하기

DataLoader에 데이터셋을 불러온 뒤에는 필요에 따라 데이터셋을 순회할 수 있다. 

이때, 각 순회마다 데이터셋은 batch_size의 크기만큼을 반환한다.

즉, 위에서는 batch_size를 64로 설정하였으므로, 각 iter(순회)마다 64개의 데이터와 64개의 label을 반환해주는 것이다.

또한, shuffle을 True로 지정했으므로, 모든 배치를 순회한 뒤에는 데이터를 섞게된다.

train_features, train_labels = next(iter(train_dataloader))
print(f"Feature batch shape : {train_features.size()} ")
print(f"Labels batch shape : {train_labels.size()}")
img = train_features[0].squeeze()
print(img.shape)
label = train_labels[0]
plt.imshow(img, cmap="gray")
plt.show()

위의 코드에서 squeeze()는 차원의 크기가 1인 차원을 제거하는 작업을 한다.

즉, train_features[0]의 크기는 1개의 데이터에만 접근하였으므로 [1, 28, 28]이 된다.

이때, squeeze()를 하여 img에 저장하였으므로, img 변수의 크기는 [28, 28]이 된다.