인스톨! 파이토치 11강 | PyTorch 커스텀 데이터셋 완벽 가이드: 나만의 데이터로 딥러닝 시작하기

[AI 인공지능 머신러닝 딥러닝/Python | PyTorch] - 인스톨! 파이토치 강의 소개

인스톨! 파이토치 강의 소개

 

인스톨! 파이토치

11-1강. 내가 가진 데이터를 딥러닝에 활용하는 법!

실전 딥러닝 프로젝트에서는 MNIST나 CIFAR-10 같은 표준 데이터셋이 아닌, 자신만의 데이터를 사용하는 경우가 대부분입니다. PyTorch에서는 torch.utils.data.Dataset 클래스를 상속받아 커스텀 데이터셋을 만들 수 있습니다.[^1][^2][^3]

Dataset 클래스의 구조

Dataset은 데이터셋을 표현하는 추상 클래스로, 커스텀 데이터셋을 만들 때는 반드시 세 가지 메소드를 구현해야 합니다:[^4][^1]

1. __init__(): 데이터 초기화 및 전처리 설정
2. __len__(): 데이터셋의 전체 샘플 개수 반환
3. __getitem__(index): 특정 인덱스의 데이터와 레이블 반환

기본 커스텀 데이터셋 구조

from torch.utils.data import Dataset
import torch

class CustomDataset(Dataset):
    def __init__(self, data, labels, transform=None):
        """
        Args:
            data: 데이터 파일 경로 리스트 또는 실제 데이터
            labels: 각 데이터에 해당하는 레이블
            transform: 데이터 전처리 함수 (optional)
        """
        self.data = data
        self.labels = labels
        self.transform = transform

    def __len__(self):
        """데이터셋의 전체 샘플 개수 반환"""
        return len(self.data)

    def __getitem__(self, idx):
        """idx번째 샘플 반환"""
        sample = self.data[idx]
        label = self.labels[idx]

        # transform이 있으면 적용
        if self.transform:
            sample = self.transform(sample)

        return sample, label

이미지 데이터셋 예제

실제로 가장 많이 사용하는 이미지 커스텀 데이터셋을 구현해보겠습니다:[^5]

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

class CustomImageDataset(Dataset):
    def __init__(self, img_dir, transform=None):
        """
        Args:
            img_dir: 이미지가 저장된 디렉토리 경로
            transform: 이미지 전처리 함수
        """
        self.img_dir = img_dir
        self.transform = transform

        # 이미지 파일 리스트와 레이블 생성
        self.img_labels = []
        self.img_paths = []

        # 클래스별 폴더 구조를 가정
        # img_dir/class1/img1.jpg
        # img_dir/class2/img2.jpg
        classes = sorted(os.listdir(img_dir))
        self.class_to_idx = {cls_name: i for i, cls_name in enumerate(classes)}

        for class_name in classes:
            class_dir = os.path.join(img_dir, class_name)
            if os.path.isdir(class_dir):
                for img_name in os.listdir(class_dir):
                    self.img_paths.append(os.path.join(class_dir, img_name))
                    self.img_labels.append(self.class_to_idx[class_name])

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

    def __getitem__(self, idx):
        img_path = self.img_paths[idx]
        image = Image.open(img_path).convert('RGB')
        label = self.img_labels[idx]

        if self.transform:
            image = self.transform(image)

        return image, label

CSV 파일 기반 데이터셋 예제

CSV 파일에 데이터와 레이블이 저장된 경우의 구현입니다:[^6]

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

class CSVImageDataset(Dataset):
    def __init__(self, csv_file, img_dir, transform=None):
        """
        Args:
            csv_file: 이미지 파일명과 레이블이 담긴 CSV 파일
            img_dir: 이미지 파일들이 있는 디렉토리
            transform: 이미지 전처리 함수
        """
        self.annotations = pd.read_csv(csv_file)
        self.img_dir = img_dir
        self.transform = transform

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

    def __getitem__(self, idx):
        # CSV에서 이미지 경로와 레이블 읽기
        img_name = os.path.join(self.img_dir, self.annotations.iloc[idx, 0])
        image = Image.open(img_name).convert('RGB')
        label = int(self.annotations.iloc[idx, 1])

        if self.transform:
            image = self.transform(image)

        return image, label

DataLoader와 함께 사용하기

커스텀 데이터셋을 만든 후 DataLoader와 함께 사용합니다:[^7]

from torch.utils.data import DataLoader
from torchvision import transforms

# Transform 정의
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])
])

# 커스텀 데이터셋 생성
dataset = CustomImageDataset(
    img_dir='./data/images',
    transform=transform
)

# DataLoader 생성
dataloader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,
    num_workers=4
)

# 사용 예시
for images, labels in dataloader:
    # 모델 학습 코드
    outputs = model(images)
    loss = criterion(outputs, labels)
    # ...

왜 커스텀 데이터셋이 필요한가?

1. 메모리 효율성: 모든 데이터를 한 번에 메모리에 로드하지 않고 필요할 때만 로드[^8]
2. 대용량 데이터 처리: RAM 용량을 초과하는 대규모 데이터셋 처리 가능[^8]
3. 유연한 전처리: 데이터 로딩 시점에 동적으로 전처리 적용
4. 다양한 데이터 형식 지원: 이미지, 텍스트, 오디오 등 모든 형식의 데이터 처리 가능[^2]

11-2강. train/test transform 함수 분리 방법!

딥러닝 모델을 학습할 때 훈련 데이터와 테스트 데이터에 서로 다른 전처리를 적용하는 것이 매우 중요합니다. 훈련 시에는 데이터 증강을 적용하여 모델의 일반화 성능을 높이지만, 테스트 시에는 원본 데이터에 가까운 상태로 평가해야 합니다.

Train과 Test Transform의 차이

Train Transform (훈련용)

train_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.RandomResizedCrop(224),      # 무작위 크롭
    transforms.RandomHorizontalFlip(p=0.5), # 무작위 좌우 반전
    transforms.RandomRotation(degrees=15),   # 무작위 회전
    transforms.ColorJitter(brightness=0.2, contrast=0.2),  # 색상 변형
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                        std=[0.229, 0.224, 0.225])
])

Test Transform (테스트용)

test_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),  # 중앙 크롭만 (무작위 X)
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                        std=[0.229, 0.224, 0.225])
])

핵심 차이점: 훈련 시에는 Random 계열 증강을 사용하지만, 테스트 시에는 일관된 전처리만 적용합니다.

방법 1: 데이터셋을 두 번 생성

가장 직관적인 방법으로, 훈련과 테스트용 데이터셋을 각각 생성합니다:[^9]

# 훈련 데이터셋
train_dataset = CustomImageDataset(
    img_dir='./data/train',
    transform=train_transform
)

# 테스트 데이터셋
test_dataset = CustomImageDataset(
    img_dir='./data/test',
    transform=test_transform
)

# DataLoader 생성
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

방법 2: train/test split 사용

하나의 데이터셋을 훈련/테스트로 분할하는 방법입니다:[^10][^11]

from torch.utils.data import random_split

# 전체 데이터셋 생성 (transform은 나중에 적용)
full_dataset = CustomImageDataset(img_dir='./data')

# 8:2 비율로 분할
train_size = int(0.8 * len(full_dataset))
test_size = len(full_dataset) - train_size

train_dataset, test_dataset = random_split(
    full_dataset, 
    [train_size, test_size],
    generator=torch.Generator().manual_seed(42)  # 재현성을 위한 시드
)

# 각각 다른 transform 적용
train_dataset.dataset.transform = train_transform
test_dataset.dataset.transform = test_transform

방법 3: 커스텀 데이터셋에 mode 파라미터 추가

더 우아한 방법으로, 데이터셋 클래스 내부에서 mode를 구분합니다:

class CustomImageDataset(Dataset):
    def __init__(self, img_dir, mode='train', transform=None):
        """
        Args:
            mode: 'train' 또는 'test'로 동작 모드 지정
        """
        self.img_dir = img_dir
        self.mode = mode

        # mode에 따라 다른 transform 설정
        if transform is None:
            if mode == 'train':
                self.transform = transforms.Compose([
                    transforms.Resize(256),
                    transforms.RandomResizedCrop(224),
                    transforms.RandomHorizontalFlip(),
                    transforms.ToTensor(),
                    transforms.Normalize([0.485, 0.456, 0.406],
                                       [0.229, 0.224, 0.225])
                ])
            else:  # test mode
                self.transform = transforms.Compose([
                    transforms.Resize(256),
                    transforms.CenterCrop(224),
                    transforms.ToTensor(),
                    transforms.Normalize([0.485, 0.456, 0.406],
                                       [0.229, 0.224, 0.225])
                ])
        else:
            self.transform = transform

        # 데이터 로딩 로직
        # ...

    def __getitem__(self, idx):
        image, label = # 이미지 로드

        if self.transform:
            image = self.transform(image)

        return image, label

# 사용 예시
train_dataset = CustomImageDataset('./data/train', mode='train')
test_dataset = CustomImageDataset('./data/test', mode='test')

방법 4: sklearn의 train_test_split 활용

sklearn을 사용하여 더 정교하게 분할할 수 있습니다:[^12]

from sklearn.model_selection import train_test_split

# 전체 데이터셋 로드
full_dataset = CustomImageDataset('./data')

# 인덱스 기반으로 분할
train_idx, test_idx = train_test_split(
    list(range(len(full_dataset))),
    test_size=0.2,
    random_state=42,
    stratify=full_dataset.labels  # 클래스 비율 유지
)

# Subset 생성
from torch.utils.data import Subset

train_dataset = Subset(full_dataset, train_idx)
test_dataset = Subset(full_dataset, test_idx)

# 각각 다른 transform wrapper 적용
class TransformDataset(Dataset):
    def __init__(self, subset, transform):
        self.subset = subset
        self.transform = transform

    def __getitem__(self, idx):
        image, label = self.subset[idx]
        if self.transform:
            image = self.transform(image)
        return image, label

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

train_dataset = TransformDataset(train_dataset, train_transform)
test_dataset = TransformDataset(test_dataset, test_transform)

실전 팁과 주의사항

1. 셔플(Shuffle) 설정

# 훈련: 셔플 O
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)

# 테스트: 셔플 X (재현성을 위해)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

2. 재현성(Reproducibility) 확보

# 랜덤 시드 고정
import random
import numpy as np

random.seed(42)
np.random.seed(42)
torch.manual_seed(42)
if torch.cuda.is_available():
    torch.cuda.manual_seed_all(42)

3. Validation set 추가

실전에서는 Train/Validation/Test 세 개로 분할합니다:

# 전체 데이터의 80% 학습, 20% 테스트
train_val_dataset, test_dataset = random_split(
    full_dataset, [0.8, 0.2]
)

# 학습 데이터의 80% 훈련, 20% 검증
train_dataset, val_dataset = random_split(
    train_val_dataset, [0.8, 0.2]
)

# Validation은 test와 동일한 transform 사용
train_dataset.transform = train_transform
val_dataset.transform = test_transform
test_dataset.transform = test_transform

4. Normalization 값 주의

ImageNet 사전 학습 모델을 사용할 때는 반드시 ImageNet 통계값을 사용해야 합니다:

mean=[0.485, 0.456, 0.406]
std=[0.229, 0.224, 0.225]

결론

커스텀 데이터셋을 만드는 것은 실전 딥러닝 프로젝트의 첫걸음입니다. Dataset 클래스를 상속받아 __init__, __len__, __getitem__을 구현하면 어떤 형식의 데이터도 PyTorch에서 사용할 수 있습니다.

훈련과 테스트 데이터에 서로 다른 transform을 적용하는 것은 모델 성능에 큰 영향을 미칩니다. 훈련 시에는 데이터 증강으로 일반화 성능을 높이고, 테스트 시에는 일관된 전처리로 공정한 평가를 수행해야 합니다. 상황에 맞는 적절한 방법을 선택하여 효율적인 데이터 파이프라인을 구축하세요.

 

 

 

다음 글

[AI 인공지능 머신러닝 딥러닝/Python | PyTorch] - 인스톨! 파이토치 12강 | PyTorch 완전한 학습 파이프라인: 실전 템플릿 코드

인스톨! 파이토치 12강 | PyTorch 완전한 학습 파이프라인: 실전 템플릿 코드