PyTorch 기초 (4) - 데이터 파이프라인: Dataset과 DataLoader

PyTorch 기초 시리즈(6편)의 4편입니다. 전체 목차는 0편에 있습니다. 모델을 학습시키는 법을 익혔으니, 이 편은 그 모델에 데이터를 어떻게 공급하는지를 다룹니다.

데이터를 통째로 넣던 방식의 한계

지금까지는 전체 데이터 X, Y를 한 번에 모델에 넣었습니다. 예제 데이터가 작아서 가능했습니다. 하지만 실제 학습에는 세 가지가 필요합니다.

• 미니배치 — 딥러닝 기초 (5)에서 봤듯, 전체가 아니라 배치 단위로 나눠 자주 갱신해야 학습이 되고 메모리도 감당됩니다.
• 셔플 — 매 에폭 데이터 순서를 섞어야 특정 순서에 모델이 편향되지 않습니다.
• 지연 로딩 — 이미지 수십만 장처럼 메모리에 다 못 올리는 데이터는 필요한 순간에 하나씩 읽어와야 합니다.

PyTorch는 이 셋을 Dataset과 DataLoader 두 부품으로 분리해 해결합니다. Dataset은 “샘플 하나를 어떻게 가져오는가”를, DataLoader는 “그것들을 어떻게 묶어 공급하는가”를 담당합니다.

Dataset: 샘플 하나를 정의한다

Dataset은 두 개의 메서드만 채우면 됩니다. 전체 개수를 알려주는 __len__과, i번째 샘플을 반환하는 __getitem__입니다.

from torch.utils.data import Dataset

class MyDataset(Dataset):
    def __init__(self, X, Y):
        self.X, self.Y = X, Y

    def __len__(self):
        return len(self.X)            # 전체 샘플 수

    def __getitem__(self, i):
        return self.X[i], self.Y[i]   # i번째 (입력, 정답)

핵심은 __getitem__이 인덱스 하나를 받아 샘플 하나를 돌려준다는 점입니다. 파일에서 이미지를 읽거나 전처리를 하는 것도 여기 안에 쓰면, 그 작업이 필요한 순간까지 미뤄집니다(지연 로딩). 이미 텐서로 준비된 데이터라면 직접 만들 것도 없이 TensorDataset을 쓰면 됩니다.

from torch.utils.data import TensorDataset

dataset = TensorDataset(X, Y)   # (X[i], Y[i])를 돌려주는 Dataset

DataLoader: 배치로 묶어 공급한다

DataLoader는 Dataset을 감싸, 샘플들을 미니배치로 묶고 섞어서 순회할 수 있게 해줍니다.

from torch.utils.data import DataLoader

loader = DataLoader(dataset, batch_size=32, shuffle=True)

for xb, yb in loader:      # xb: (32, ...), yb: (32, ...)
    ...                    # 배치 하나씩 꺼내진다

• batch_size — 한 번에 묶을 샘플 수. DataLoader가 __getitem__으로 32개를 꺼내 자동으로 쌓아(batch) 줍니다.
• shuffle=True — 매 에폭 순서를 섞습니다. 학습에는 켜고, 평가에는 보통 끕니다.
• num_workers — 데이터를 여러 프로세스로 병렬 로딩. 파일 읽기·전처리가 무거우면 학습 속도가 붙습니다.

학습 루프에 끼워 넣기

(3)편의 학습 루프에서 달라지는 것은 딱 한 줄입니다. 전체를 한 번에 넣던 자리가, 배치를 순회하는 안쪽 루프로 바뀝니다.

for epoch in range(100):
    for xb, yb in loader:          # ← 배치 순회가 추가된다
        pred = model(xb)
        loss = criterion(pred, yb)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

안쪽의 다섯 줄(순전파 → 손실 → zero_grad → backward → step)은 그대로입니다. 한 에폭이 이제 “전체를 한 번 갱신”이 아니라 “배치마다 한 번씩, 데이터 전체를 한 바퀴 갱신”이 됩니다. 이미지·텍스트 같은 데이터의 변형(증강)이나 정규화는 torchvision.transforms 등을 Dataset에 붙여 처리하는데, 이는 실습편에서 실제 데이터와 함께 다룹니다.

정리

개념	한 줄 요약
Dataset	샘플 하나를 정의. __len__ + __getitem__(i)
__getitem__	인덱스 → 샘플 하나. 파일 읽기·전처리를 지연 실행
TensorDataset	이미 텐서인 데이터를 감싸는 간편 Dataset
DataLoader	배치로 묶고(batch_size) 섞어(shuffle) 공급
num_workers	여러 프로세스로 병렬 로딩
학습 루프	에폭 안에 for xb, yb in loader 배치 순회 추가

이제 모델·학습·데이터 공급이 모두 갖춰졌습니다. 남은 것은 이걸 실제로 GPU에서 안정적으로, 재현 가능하게 돌리는 실무 감각입니다. 마지막 편에서 정리합니다.

다음 글: PyTorch 기초 (5) - 학습 실무: GPU, train/eval 모드, 체크포인트, 재현성