[PYTHON] 커스텀 데이터셋(Dataset) 클래스 제작 방법 7가지와 데이터 파이프라인 해결 전략

커스텀 데이터셋 (Dataset) 클래스

 

딥러닝 모델의 성능을 결정짓는 핵심 요소는 모델의 아키텍처보다 '데이터의 품질과 이를 처리하는 파이프라인의 효율성'에 있습니다. PyTorch나 TensorFlow 프레임워크에서 제공하는 기본 데이터셋 외에, 기업의 내부 데이터나 비정형 데이터를 모델이 학습할 수 있는 형태로 가공하려면 커스텀 데이터셋(Custom Dataset) 클래스 구현 능력이 필수입니다. 본 가이드에서는 파이썬을 활용해 다양한 형태의 데이터를 다루는 7가지 실무 예제와 함께, 메모리 병목 현상을 해결하는 전문적인 데이터 파이프라인 구축 노하우를 상세히 다룹니다.

---

1. 커스텀 데이터셋 클래스의 필수 구성 요소와 원리

PyTorch의 torch.utils.data.Dataset을 상속받아 커스텀 클래스를 만들 때, 반드시 구현해야 하는 3가지 핵심 메서드가 있습니다. 이 구조를 이해하는 것이 모든 데이터 파이프라인 설계의 시작입니다.

메서드 명칭	역할 및 기능	구현 필수 여부
__init__	데이터의 경로, 라벨링 파일 읽기, 데이터 전처리 정의	필수 (Initial Setup)
__len__	전체 데이터셋의 크기(샘플 개수)를 반환	필수 (Dataset Size)
__getitem__	특정 인덱스의 샘플을 읽어와 텐서로 변환하여 반환	필수 (Data Retrieval)
_preprocess	이미지 크기 조정, 텍스트 토큰화 등 내부 로직 처리	선택 (Internal Logic)

---

2. 실무자를 위한 커스텀 데이터셋 구현 예제 7가지

이미지, 텍스트, CSV, 오디오 등 실무에서 가장 많이 접하는 데이터 형식을 처리하는 해결 방법을 코드 위주로 살펴봅니다.

Example 1: 이미지 폴더 구조를 읽어오는 기본 커스텀 데이터셋

파일 경로와 라벨 정보를 리스트로 관리하여 이미지를 로드하는 표준 방식입니다.

import os
import torch
from torch.utils.data import Dataset
from PIL import Image

class ImageCustomDataset(Dataset):
    def __init__(self, image_dir, labels, transform=None):
        self.image_dir = image_dir
        self.labels = labels
        self.image_files = os.listdir(image_dir)
        self.transform = transform

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

    def __getitem__(self, idx):
        img_path = os.path.join(self.image_dir, self.image_files[idx])
        image = Image.open(img_path).convert("RGB")
        label = torch.tensor(self.labels[idx])
        
        if self.transform:
            image = self.transform(image)
        return image, label

print("기본 이미지 커스텀 데이터셋 클래스 정의 완료")
    

Example 2: CSV 파일 기반의 정형 데이터셋 구현

Pandas를 활용하여 엑셀이나 CSV에 저장된 수치 데이터를 학습용 텐서로 변환하는 해결책입니다.

import pandas as pd

class CsvDataset(Dataset):
    def __init__(self, csv_file):
        self.data = pd.read_csv(csv_file)
        # 특징 데이터와 타겟 데이터를 분리
        self.x = self.data.iloc[:, :-1].values
        self.y = self.data.iloc[:, -1].values

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

    def __getitem__(self, idx):
        features = torch.FloatTensor(self.x[idx])
        target = torch.LongTensor([self.y[idx]])
        return features, target
    

Example 3: NLP 실무를 위한 텍스트 토큰화 데이터셋

문장을 입력받아 토크나이저를 거쳐 패딩 처리가 된 정수 인덱스로 반환하는 구조입니다.

class TextDataset(Dataset):
    def __init__(self, texts, labels, tokenizer, max_len=128):
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer
        self.max_len = max_len

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

    def __getitem__(self, idx):
        text = str(self.texts[idx])
        encoding = self.tokenizer.encode_plus(
            text, add_special_tokens=True, max_length=self.max_len,
            padding='max_length', truncation=True, return_tensors='pt'
        )
        return {
            'input_ids': encoding['input_ids'].flatten(),
            'attention_mask': encoding['attention_mask'].flatten(),
            'labels': torch.tensor(self.labels[idx], dtype=torch.long)
        }
    

Example 4: 시계열(Time-Series) 윈도우 슬라이싱 데이터셋

연속된 수치 데이터에서 특정 길의의 윈도우를 추출하여 다음 시점을 예측하는 해결 방법입니다.

class TimeSeriesDataset(Dataset):
    def __init__(self, data, window_size=10):
        self.data = torch.FloatTensor(data)
        self.window_size = window_size

    def __len__(self):
        return len(self.data) - self.window_size

    def __getitem__(self, idx):
        x = self.data[idx:idx + self.window_size]
        y = self.data[idx + self.window_size]
        return x, y
    

Example 5: 데이터 증강(Augmentation)이 포함된 동적 데이터셋

학습 시마다 매번 다른 변형을 가하여 모델의 일반화 성능을 높이는 실무 구현법입니다.

from torchvision import transforms

train_transform = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(15),
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

# 위에서 정의한 ImageCustomDataset에 train_transform을 인자로 전달하여 사용
    

Example 6: 대용량 데이터 로딩을 위한 Lazy Loading 기법

모든 데이터를 메모리에 올리지 않고, 요청 시에만 디스크에서 읽어와 OOM(메모리 부족)을 해결하는 전략입니다.

class LazyDataset(Dataset):
    def __init__(self, file_paths):
        self.file_paths = file_paths # 경로 리스트만 보관

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

    def __getitem__(self, idx):
        # 호출 시점에만 파일을 읽음
        data = torch.load(self.file_paths[idx])
        return data
    

Example 7: 다중 라벨(Multi-label) 분류를 위한 데이터셋

하나의 데이터에 여러 개의 타겟값이 할당된 경우를 처리하는 해결 방식입니다.

class MultiLabelDataset(Dataset):
    def __init__(self, features, labels_list):
        self.features = features
        self.labels = labels_list # [[1, 0, 1], [0, 1, 1]...] 형태

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

    def __getitem__(self, idx):
        x = torch.FloatTensor(self.features[idx])
        y = torch.FloatTensor(self.labels[idx])
        return x, y
    

---

3. 데이터 파이프라인의 성능을 높이는 3가지 차이점 해결책

단순히 클래스를 만드는 것을 넘어, GPU 연산 효율을 극대화하기 위한 튜닝 전략입니다.

1. NumWorkers 설정: DataLoader의 num_workers를 CPU 코어 수에 맞게 설정하여 멀티 프로세싱으로 데이터를 로드하세요.
2. Pin Memory 활용: pin_memory=True를 설정하면 CPU에서 GPU로의 데이터 전송 속도가 획기적으로 향상됩니다.
3. Prefetching: 데이터 로딩과 모델 연산을 병렬화하여 GPU가 대기하는 시간을 줄여야 합니다.

---

4. 결론 및 요약

커스텀 데이터셋 제작은 단순히 파이썬 클래스를 상속받는 행위가 아니라, 데이터의 흐름(Flow)을 설계하는 과정입니다. 정형 데이터부터 비정형 멀티모달 데이터까지, 상황에 맞는 __getitem__ 로직을 구현하고 DataLoader와의 유기적인 연결을 통해 효율적인 딥러닝 학습 환경을 구축하시기 바랍니다.

 

내용 출처 및 참조:

• PyTorch Official Documentation: torch.utils.data.Dataset
• Deep Learning with PyTorch by Eli Stevens, Luca Antiga.
• NVIDIA DALI: Fast Data Loading for Deep Learning.
• Hugging Face Datasets Library Guide.