PyTorch 효율적인 학습을 위한 Dataset과 DataLoader 활용법 (Custom 예제 포함)

딥러닝 모델 학습의 절반은 데이터 준비입니다. 아무리 훌륭한 모델 아키텍처를 설계했더라도, 데이터를 모델에 공급하는 과정에서 병목이 발생하면 학습 효율은 급격히 떨어집니다.

특히 메모리 제한이 있는 임베디드 환경이나 대규모 GPU 서버에서 데이터를 효과적으로 관리하려면 PyTorch의 핵심 도구인 Dataset과 DataLoader를 제대로 이해하고 활용해야 합니다. 본 포스팅에서는 이 두 클래스의 개념과 실무적인 커스텀 활용법을 정리해 보겠습니다.

---

핵심 요약

1. Dataset: 데이터셋을 로드하고 개별 샘플에 대한 전처리를 담당하는 클래스입니다.
2. DataLoader: Dataset을 감싸서 배치(Batch) 생성, 셔플(Shuffle), 병렬 로딩을 수행하는 엔진입니다.
3. 최적화: num_workers와 pin_memory 설정을 통해 데이터 로딩 병목 현상을 해결할 수 있습니다.

---

Generated by Gemini AI

1. Dataset 클래스: 데이터의 정의

Dataset은 "무엇을, 어떻게 가져올 것인가"를 정의하는 단계입니다. PyTorch의 torch.utils.data.Dataset을 상속받아 커스텀 클래스를 만들면, CSV, 이미지, 오디오 등 어떤 형태의 데이터라도 모델이 이해할 수 있는 표준 텐서로 가공할 수 있습니다.

커스텀 Dataset 구현 시 필수 메서드

• __init__: 데이터의 경로를 설정하거나 파일 목록을 불러오는 초기화 단계입니다.
• __len__: 전체 데이터셋의 크기를 반환합니다.
• __getitem__: 특정 인덱스(idx)의 데이터를 가져와 전처리 후 텐서로 반환합니다.

Python

 

import torch
from torch.utils.data import Dataset
import pandas as pd

class CustomDataset(Dataset):
    def __init__(self, csv_file):
        # 실무 Tip: 데이터가 너무 크면 여기서 읽지 않고 경로 리스트만 만듭니다.
        self.data = pd.read_csv(csv_file)
        
    def __len__(self):
        return len(self.data)
        
    def __getitem__(self, idx):
        # 데이터를 불러와 텐서로 변환
        sample = torch.tensor(self.data.iloc[idx].values, dtype=torch.float32)
        return sample

---

2. DataLoader: 데이터 공급의 엔진

Dataset이 식재료라면, DataLoader는 이를 요리하기 좋게 '배차'하는 역할을 합니다. 단순히 데이터를 가져오는 것을 넘어, 모델 학습을 위한 고성능 파이프라인을 구축해 줍니다.

주요 파라미터와 최적화 팁

파라미터	설명	실무 최적화 팁
batch_size	한 번에 로드할 데이터 개수	GPU VRAM 용량에 맞춰 2의 거듭제곱으로 설정 권장
shuffle	데이터를 랜덤하게 섞을지 여부	학습(Train) 시에는 필수, 검증 시에는 False
num_workers	병렬 처리에 사용할 프로세스 수	CPU 코어 수의 1/2~전부 사이에서 최적점을 찾으세요
pin_memory	GPU 전송 가속 여부	GPU 학습 시 True로 설정하면 전송 속도가 빨라집니다

Python

 

from torch.utils.data import DataLoader

dataset = CustomDataset("data.csv")
dataloader = DataLoader(
    dataset, 
    batch_size=32, 
    shuffle=True, 
    num_workers=4, 
    pin_memory=True
)

# 학습 루프 예시
for batch in dataloader:
    # batch: [32, feature_dim] 형태의 텐서가 자동 공급됨
    optimizer.zero_grad()
    output = model(batch)
    ...

---

3. 실전 활용: 이미지 데이터 처리와 Transform

이미지 데이터를 다룰 때는 torchvision.transforms와 결합하여 데이터 증강(Augmentation)과 정규화를 동시에 수행하는 것이 효율적입니다.

Python

 

from torchvision import transforms
from PIL import Image
import os

class ImageDataset(Dataset):
    def __init__(self, image_dir, transform=None):
        self.image_dir = image_dir
        # 실무 Tip: 이미지 파일 확장자만 필터링하여 에러를 방지합니다.
        self.image_filenames = [f for f in os.listdir(image_dir) if f.endswith(('.png', '.jpg'))]
        self.transform = transform
        
    def __len__(self):
        return len(self.image_filenames)
        
    def __getitem__(self, idx):
        img_path = os.path.join(self.image_dir, self.image_filenames[idx])
        image = Image.open(img_path).convert("RGB")
        
        if self.transform:
            image = self.transform(image)
        return image

# 전처리 파이프라인 구성 예시
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

---

4. 주의할 점 및 문제 해결 (Troubleshooting)

1. 데이터 로딩 병목: GPU 사용률이 낮다면 num_workers를 높여보세요. CPU가 데이터를 준비하는 속도가 GPU 연산 속도를 못 따라가는 경우가 많습니다.
2. Out Of Memory (OOM): batch_size가 너무 크면 발생합니다. 에러 발생 시 배치 사이즈를 절반으로 줄여보세요.
3. 학습 재현성: shuffle=True로 설정하더라도, 실험의 재현성을 위해 torch.manual_seed()를 사용하여 난수 시드를 고정하는 습관이 중요합니다.

---

결론

PyTorch의 Dataset과 DataLoader는 단순한 데이터 로더가 아니라, 복잡한 데이터를 정제하여 모델에 끊김 없이 공급하는 엔지니어링 파이프라인입니다.

• Dataset으로 데이터의 개별 가공 로직을 캡슐화하세요.
• DataLoader의 파라미터를 튜닝하여 하드웨어 성능을 최대로 끌어내세요.

오늘 소개한 예제 코드를 여러분의 프로젝트에 맞춰 수정해 사용해 보시기 바랍니다. 데이터 파이프라인이 탄탄할수록 모델의 연구 속도는 비약적으로 빨라질 것입니다.