Pytorch에서의 Dataset, DataLoader

Pytorch에서는 torch.utils.data.DataLoader, torch.utils.data.Dataset 이 2개를 통해 준비된 데이터셋, 가지고 있는 데이터셋을 모두 사용할 수 있다.

Dataset은 샘플과 label을 저장하고, DataLoader은 Dataset을 샘플에 쉽게 접근할 수 있도록 순회 가능한 객체로 감싼다.

미리 정의된 데이터셋 불러오기 - FashionMnist

FashionMnist 데이터는 60000개의 학습 데이터, 10000개의 테스트 데이터로 이루어져 있다. 각각 흑백 28x28 이미지, 10개의 label 중 하나의 정답 데이터로 구성된다.

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.FashionMNIST(
    root="data", # 데이터 저장 경로
    train=True, # Train, Test 전용 데이터 결정
    download=True, # root에 없을 때 데이터 다운로드 여부
    transform=ToTensor() # feature와 label transform 지정
)

test_data = datasets.FashionMNIST(
    root="data",
    train=False,
    download=True,
    transform=ToTensor()
)

데이터셋 순회, 시각화

Dataset에 리스트처럼 직접 접근할 수 있다.

labels_map = {
    0: "T-Shirt",
    1: "Trouser",
    2: "Pullover",
    3: "Dress",
    4: "Coat",
    5: "Sandal",
    6: "Shirt",
    7: "Sneaker",
    8: "Bag",
    9: "Ankle Boot",
}
figure = plt.figure(figsize=(8, 8))
cols, rows = 3, 3
for i in range(1, cols * rows + 1):
    sample_idx = torch.randint(len(training_data), size=(1,)).item()
    img, label = training_data[sample_idx]
    figure.add_subplot(rows, cols, i)
    plt.title(labels_map[label])
    plt.axis("off")
    plt.imshow(img.squeeze(), cmap="gray")
plt.show()

Custom Dataset 만들기

Custom Dataset은 3개 함수를 반드시 구현해야 한다.

• init
• len
• getitem

우선 Fashion MNIST 데이터를 사용하지만, imgdir에 이미지가 저장되고, annotation_file에 정답이 별도로 저장된다.

import os
import pandas as pd
from torchvision.io import read_image

class CustomImageDataset(Dataset):
    def __init__(self, annotations_file, img_dir, transform=None, target_transform=None):
        self.img_labels = pd.read_csv(annotations_file, names=['file_name', 'label'])
        self.img_dir = img_dir
        self.transform = transform
        self.target_transform = target_transform

    def __len__(self):
        return len(self.img_labels)

    def __getitem__(self, idx):
        img_path = os.path.join(self.img_dir, self.img_labels.iloc[idx, 0])
        image = read_image(img_path)
        label = self.img_labels.iloc[idx, 1]
        if self.transform:
            image = self.transform(image)
        if self.target_transform:
            label = self.target_transform(label)
        return image, label

__init__

def __init__(self, annotations_file, img_dir, transform=None, target_transform=None):
    self.img_labels = pd.read_csv(annotations_file) # annotations_file을 csv로 불러온다.
    self.img_dir = img_dir
    self.transform = transform
    self.target_transform = target_transform

해당 함수에서는 Dataset 객체가 생성될 때 한 번만 실행된다. 여기서는 이미지와 annotation_file이 포함된 디렉토리와 두가지 변형을 초기화한다.

labels.csv는 다음과 같이 구성된다.

tshirt1.jpg, 0
tshirt2.jpg, 0
......
ankleboot999.jpg, 9

__len__

데이터셋의 개수를 반환한다.

def __len__(self):
    return len(self.img_labels)

__getitem__

주어진 인덱스 idx에 해당하는 샘플을 데이터셋에서 불러오고 반환하는 과정을 수행한다. 인덱스를 기반으로, 디스크에서 이미지의 위치를 식별하고, read_image를 사용해 이미지 텐서로 변환, self.img_labels의 csv 데이터로부터 해당하는 label을 가져오고, 만약 해당한다면 변형 함수들을 호출한 뒤, 텐서 이미지와 라벨을 dictionary 형태로 반환한다.

def __getitem__(self, idx):
    img_path = os.path.join(self.img_dir, self.img_labels.iloc[idx, 0])
    image = read_image(img_path)
    label = self.img_labels.iloc[idx, 1]
    if self.transform:
        image = self.transform(image)
    if self.target_transform:
        label = self.target_transform(label)
    sample = {"image": image, "label": label}
    return sample

DataLoader로 학습용 데이터 준비

Dataset은 데이터셋의 feature을 가져오고 하나의 샘플에 label을 지정하는 일을 한 번에 수행할 수 있다. 모델 학습 시, 일반적인 샘플들을 미니배치 형식으로 전달하고, 매 에포크마다 데이터를 다시 섞어 Overfit을 막는다. 추가로, multiprocessing을 사용하면 데이터 검색 속도를 높일 수 있다.

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)

DataLoader(dataset, batch_size=1, shuffle=False, sampler=None,
           batch_sampler=None, num_workers=0, collate_fn=None,
           pin_memory=False, drop_last=False, timeout=0,
           worker_init_fn=None, *, prefetch_factor=2,
           persistent_workers=False)

DataLoader을 통해 순회하기(iterate)

DataLoader에 데이터셋을 불러온 뒤, 필요에 따라 데이터셋을 순회할 수 있다. 아래의 각 순회 과정은 train_features, train_labels의 묶음(batch)를 반환한다. 위에서 shuffle=True로 지정했기에, 모든 배치를 순회한 뒤 데이터가 섞인다.

# 이미지와 정답(label)을 표시합니다.
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()
label = train_labels[0]
plt.imshow(img, cmap="gray")
plt.show()
print(f"Label: {label}")

Transform

딥러닝을 학습할 때나 머신러닝을 학습할 때, 항상 최종 처리가 된 형태로 제공하지는 않는다. 그렇기에 우리는 transform을 통해 데이터를 조작하고 학습에 적합하게 만드는 과정이 필요하다.

여기서는 torchvision을 중심으로 사용한다.

모든 TorchVision 데이터셋은 변형 로직을 갖는, 호출 가능한 객체를 받는 파라미터 2개(transform, target_transform)를 갖는다.

FashionMNIST의 경우, feature는 PIL Image, label은 Integer 형태를 가진다. 학습을 위해서는 정규화된 텐서의 feature과 one-hot 형태로 인코딩 된 레이블이 필요하다.

import torch
from torchvision import datasets
from torchvision.transforms import ToTensor, Lambda

ds = datasets.FashionMNIST(
    root="data",
    train=True,
    download=True,
    transform=ToTensor(),
    target_transform=Lambda(lambda y: torch.zeros(10, dtype=torch.float).scatter_(0, torch.tensor(y), value=1))
)

여기서 ToTensor()은 PIL Image나 Numpy ndarray를 FloatTensor로 변환하고, 이미지의 픽셀의 크기를 0~1 범위로 비례해 조정한다.

Lambda Transform

Lambda 코드를 사용해 정수를 one-hot 형태로 바꾸는 역할을 수행한다.

크기가 10인 zero tensor을 새성하고, scatter_을 통해 주어진 정답에 해당하는 인덱스를 할당한다.

target_transform = Lambda(lambda y: torch.zeros(
    10, dtype=torch.float).scatter_(dim=0, index=torch.tensor(y), value=1))

참조

https://tutorials.pytorch.kr/beginner/basics/transforms_tutorial.html