DataLoader와 배치 처리

DataLoader가 맡는 일

DataLoader는 Dataset을 감싸서 여러 기능을 더해 줘. N개 예시를 [N, ...] 형태의 텐서로 묶는 배치 처리, 샘플 순서를 바꾸는 셔플, 여러 워커 프로세스를 활용한 병렬 로딩, CPU→GPU 전송을 빠르게 하는 pin 메모리가 대표적이야. 학습 루프 맨 위에서 한 줄로 설정하면, 나머지 코드는 배치를 반복해서 처리하기만 하면 돼.

가장 중요한 인자는 두 가지야. batch_size는 VRAM과 처리량에 맞춰 조정하고, num_workers는 CPU와 입출력 대역폭에 맞춰 조정해. 기본값인 num_workers=0은 단일 프로세스로 동작해. 작은 데이터셋에는 괜찮지만, 그 밖의 경우에는 디버깅 악몽이 될 수 있어.

팁: num_workers는 사용 가능한 물리 CPU 코어 수에 맞춰 잡아. 16 core 장비라면 num_workers=8이 합리적인 시작점이야. htop을 열어 두고 테스트해 봐. 워커와 GPU가 모두 바쁘게 움직이고, 어느 한쪽도 다른 쪽을 기다리며 놀지 않아야 해.

자주 빠지는 함정

학습에서 shuffle=True를 빼먹는 경우 — 데이터가 클래스별로 정렬되어 있으면 모델이 첫 번째 클래스에 과적합한 뒤 회복하지 못해. 학습 로더에서는 항상 셔플해야 해.

검증 로더를 셔플하는 경우 — 시간만 낭비하고 배치별 지표의 재현성도 깨뜨려. 검증에서는 shuffle=False로 설정해.

CUDA에서 num_workers>0을 쓰면서 pin_memory=True를 빠뜨리는 경우 — CPU→GPU 대역폭를 낭비하게 돼. 둘은 항상 함께 설정해.

__getitem__에서 공유 상태를 변경하는 경우 — 워커는 subprocess이므로 전역 상태도 워커마다 따로 존재해. 기록 카운터와 RNG 상태 등이 예상과 미묘하게 다르게 동작할 수 있어.

맞춤형 collate

기본 collate_fn은 형태가 같은 텐서를 쌓아 줘. 길이가 제각각인 시퀀스인 텍스트, 오디오, 세트를 다룰 때는 pad나 pack을 수행하는 맞춤형 collate가 필요해. Hugging Face의 DataCollatorWithPadding은 흔히 만나는 경우를 처리해 줘.

원칙: 모든 학습 DataLoader에는 세 가지를 설정해. shuffle=True, num_workers>0, 그리고 CUDA를 쓴다면 pin_memory=True야. 검증에도 세 가지를 설정해. shuffle=False, num_workers>0, pin_memory=True야.

Code

올바른 학습·검증 DataLoader 구성·python

from torch.utils.data import DataLoader

train_loader = DataLoader(
    train_ds,
    batch_size=128,
    shuffle=True,
    num_workers=8,
    pin_memory=True,
    persistent_workers=True,   # avoid worker spawn cost each epoch
    drop_last=True,            # drop the last partial batch (cleaner stats)
)

val_loader = DataLoader(
    val_ds,
    batch_size=256,            # bigger batch for inference
    shuffle=False,
    num_workers=4,
    pin_memory=True,
    persistent_workers=True,
)

가변 길이 시퀀스용 맞춤형 collate·python

import torch
from torch.nn.utils.rnn import pad_sequence

def pad_collate(batch):
    seqs, labels = zip(*batch)        # tuples of variable-length tensors and scalars
    seqs = pad_sequence(seqs, batch_first=True, padding_value=0)
    labels = torch.tensor(labels)
    return seqs, labels

loader = DataLoader(my_dataset, batch_size=32, collate_fn=pad_collate)

Exercise

직접 만든 맞춤형 데이터셋에 DataLoader를 적용해 봐. 한 에폭은 num_workers=0으로, 다른 한 에폭은 num_workers=8, pin_memory=True로 실행 시간을 재고 속도가 얼마나 빨라졌는지 확인해. num_workers=8에서도 GPU가 데이터를 기다린다면 병목은 연산이 아니라 데이터 로딩에 있어.

DataLoader가 맡는 일

자주 빠지는 함정

맞춤형 collate

Code

External links

Exercise

Progress

댓글 0