재미있는 삽질 기록

Dataset과 DataLoader, 그리고 데이터 파이프라인

#pytorch#dataset#dataloader

PyTorch에서 데이터가 모델에 들어가기까지의 파이프라인을 Dataset, DataLoader, 그리고 이미지 변환인 ToTensor을 중심으로 정리.

Dataset 클래스#

PyTorch에서는 데이터의 전처리와 배치 처리를 용이하게 할 수 있도록 Dataset과 DataLoader 클래스를 사용한다.

Dataset 클래스는 데이터셋을 정의하는 기본 클래스로서, Dataset을 상속받아 사용자 정의 데이터셋을 만든다. 보통 세 개의 메서드로 구성된다.

  • __init__: 데이터를 초기화하는 메서드
  • __len__: 데이터의 길이를 반환하는 메서드
  • __getitem__: 특정 인덱스의 데이터 샘플을 반환하는 메서드
from torch.utils.data import Dataset
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split

class IrisDataset(Dataset):
    def __init__(self, mode="train", random_state=0):
        iris = load_iris()
        train_X, test_X, train_y, test_y = train_test_split(
            iris.data, iris.target,
            stratify=iris.target, test_size=0.2,
            random_state=random_state,
        )
        if mode == "train":
            self.X = torch.tensor(train_X, dtype=torch.float32)
            self.y = torch.tensor(train_y, dtype=torch.int64)
        else:
            self.X = torch.tensor(test_X, dtype=torch.float32)
            self.y = torch.tensor(test_y, dtype=torch.int64)

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

    def __getitem__(self, idx):
        return self.X[idx], self.y[idx]

여기서 feature X는 float32, 목표 라벨 y는 int64 또는 long 으로 변환한다. 분류의 정답 라벨은 정수형이어야 하기 때문이다. 데이터의 dtype은 항상 중요하다.

DataLoader 클래스#

DataLoader 클래스는 Dataset의 인스턴스를 감싸서 배치 단위로 데이터를 로드하고, 데이터셋을 셔플링하는 등의 작업을 수행한다.

from torch.utils.data import DataLoader

batch_size = 16

train_dataset = IrisDataset("train", random_state=42)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

test_dataset = IrisDataset("test", random_state=42)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

shuffle 설정이 train과 test에서 다른 점이 중요하다.

  • 모델 훈련 시에는 데이터의 순서에 따른 편향을 줄이기 위해 데이터를 섞는다. (shuffle=True)
  • 모델 성능 평가 시에는 데이터 순서를 유지하는 것이 일반적이므로 데이터를 섞지 않는다. (shuffle=False)

train data는 셔플하고 test data는 셔플하지 않는다. 참고로 DataLoader의 shuffle은 한 epoch이 완료될 때마다 데이터를 다시 섞는다.

이미지 데이터와 ToTensor#

이미지 데이터를 다룰 때는 한 단계가 더 있다. 이미지 타입에는 PIL과 OpenCV가 있는데, PyTorch는 PIL 타입의 image를 선호한다.

문제는 DataLoader가 PIL 데이터를 그대로 배치화하지 못한다는 점이다. 그래서 torchvision.transforms.ToTensor() 클래스로 PIL을 tensor로 변환한 후 DataLoader에 넣어야 한다. ToTensor는 두 가지 일을 한다.

  1. PIL의 (H, W, C)를 tensor의 (C, H, W)로 바꾼다. (Channel Last → Channel First)
  2. 데이터를 0~1 사이의 float 값으로 정규화한다.
from torchvision import transforms
from torchvision.datasets import MNIST

mnist_train = MNIST(root=".", train=True,  download=True, transform=transforms.ToTensor())
mnist_test  = MNIST(root=".", train=False, download=True, transform=transforms.ToTensor())

DataLoader 내부를 눈으로 확인하기#

DataLoader가 실제로 어떤 배치를 내놓는지 확인하고 싶으면 next(iter(...))를 쓰면 된다.

x_batch, y_batch = next(iter(train_loader))

iter로 반복자를 만들고 next로 첫 배치를 꺼내, batch 모양과 값을 직접 찍어볼 수 있다.

댓글