[PyTorch] Dataset & DataLoader with CIFAR-10

⚠️ 본 내용은 PyTorch Korea의 공식 문서에 기반하여 공부한 내용을 적은것이니 양해바랍니다!

Dataset과 DataLoader

---

Dataset & DataLoader

PyTorch의 Dataset과 DataLoader는 데이터셋을 효율적으로 로드하고 전처리하기 위해 사용되는 강력한 도구입니다.
대규모 데이터셋을 배치 단위로 나누어 모델에 공급할 수 있도록 도와줍니다.

 

• 데이터 샘플을 처리하는 코드는 지저분(messy)하고 유지보수가 어려울 수 있습니다.
• 더 나은 가독성(readability)과 모듈성(modularity)을 위해 데이터셋 코드를 모델 학습 코드로부터 분리하는 것이 이상적입니다.
• PyTorch는 torch.utils.data.DataLoader 와 torch.utils.data.Dataset 의 두 가지 데이터 기본 요소를 제공하여 미리 준비해둔(pre-loaded) 데이터셋 뿐만 아니라 가지고 있는 데이터를 사용할 수 있도록 합니다.
• Dataset은 샘플과 정답(label)을 저장하고, DataLoader 는 Dataset 을 샘플에 쉽게 접근할 수 있도록 순회 가능한 객체(iterable)로 감쌉니다.

 

• PyTorch의 도메인 특화 라이브러리들은 (FashionMNIST와 같은) 미리 준비해둔(pre-loaded) 다양한 데이터셋을 제공합니다.
• 데이터셋은 torch.utils.data.Dataset 의 하위 클래스로 개별 데이터를 특정하는 함수가 구현되어 있습니다.
• 이러한 데이터셋은 모델을 만들어보고(prototype) 성능을 측정(benchmark)하는데 사용할 수 있습니다.
• 아래는 여러 데이터셋 예시입니다. 링크를 달아놓았으니 확인해보세요!

Image Dataset

Datasets — Torchvision 0.19 documentation

Text Dataset

torchtext.datasets — Torchtext 0.18.0 documentation

Audio Dataset

torchaudio.datasets — Torchaudio 2.4.0 documentation

---

Dataset 불러오기

TorchVision에서 CIFAR-10 데이터셋을 불러오는 예제를 살펴보겠습니다.
공식문서에는 MNIST로 되어 있지만, 따라하는거 같아서 한번 다른 데이터셋으로 전 시도해 보겠습니다!

CIFAR-10 and CIFAR-100 datasets

• CIFAR-10 데이터셋은 60,000개의 32x32 컬러 이미지로 구성되어 있으며, 각 이미지는 10개의 클래스 중 하나에 속합니다.
• 50,000개의 학습 이미지와 10,000개의 테스트 이미지로 나누어져 있습니다.

CIFAR-10 Dataset Parameter (매개변수)

• root 는 학습/테스트 데이터가 저장되는 경로입니다.
• train 은 학습용 또는 테스트용 데이터셋 여부를 지정합니다.
• download=True 는 root 에 데이터가 없는 경우 인터넷에서 다운로드합니다.
• transform 과 target_transform 은 특징(feature)과 정답(label) 변형(transform)을 지정합니다.

import torch
from torchvision import datasets, transforms
import matplotlib.pyplot as plt

# 데이터 변환 정의 (정규화 포함)
transform = transforms.Compose([
    transforms.ToTensor(),  # 이미지를 텐서로 변환
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # 이미지 정규화
])

# CIFAR-10 학습 데이터셋 불러오기
training_data = datasets.CIFAR10(
    root="data",
    train=True,
    download=True,
    transform=transform
)

# CIFAR-10 테스트 데이터셋 불러오기
test_data = datasets.CIFAR10(
    root="data",
    train=False,
    download=True,
    transform=transform
)

---

Dataset을 순회하고 시각화 하기

Dataset에 리스트(list)처럼 직접 접근(index)할 수 있습니다.
training_data[index]. matplotlib 을 사용하여 학습 데이터의 일부를 시각화해보겠습니다.

# CIFAR-10 클래스 레이블 맵
labels_map = {
    0: "Airplane",
    1: "Automobile",
    2: "Bird",
    3: "Cat",
    4: "Deer",
    5: "Dog",
    6: "Frog",
    7: "Horse",
    8: "Ship",
    9: "Truck"
}

# 시각화
figure = plt.figure(figsize=(8, 8))
cols, rows = 3, 3
for i in range(1, cols * rows + 1):
    sample_idx = torch.randint(len(training_data), size=(1,)).item()
    img, label = training_data[sample_idx]
    figure.add_subplot(rows, cols, i)
    plt.title(labels_map[label])
    plt.axis("off")
    img = img / 2 + 0.5  # 정규화 해제
    plt.imshow(img.permute(1, 2, 0))  # (C, H, W) -> (H, W, C)로 변경
plt.show()

---

파일에서 사용자 정의 데이터셋 만들기

• 사용자 정의 Dataset 클래스는 반드시 3개 함수를 구현해야 합니다.
• __init__, __len__, and __getitem__. 아래 구현을 살펴보면 CIFAR-10 이미지들은 img_dir 디렉토리에 저장되고, 정답은 annotations_file csv 파일에 별도로 저장됩니다.
• 한번 각 함수들에서 일어나는 일들을 자세히 살펴보겠습니다.

import os
import pandas as pd
from torchvision.io import read_image
from torch.utils.data import Dataset

class CustomImageDataset(Dataset):
    def __init__(self, annotations_file, img_dir, transform=None, target_transform=None):
        self.img_labels = pd.read_csv(annotations_file)
        self.img_dir = img_dir
        self.transform = transform
        self.target_transform = target_transform

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

    def __getitem__(self, idx):
        img_path = os.path.join(self.img_dir, self.img_labels.iloc[idx, 0])
        image = read_image(img_path)
        label = self.img_labels.iloc[idx, 1]
        if self.transform:
            image = self.transform(image)
        if self.target_transform:
            label = self.target_transform(label)
        sample = {"image": image, "label": label}
        return sample

__init__

• __init__ 함수는 Dataset 객체가 생성(instantiate)될 때 한 번만 실행됩니다.
• 여기서는 이미지와 주석 파일(annotation_file)이 포함된 디렉토리와 (다음 장에서 자세히 살펴볼) 두가지 변형(transform)을 초기화합니다. labels.csv 파일은 예시를 들면 다음과 같습니다.

tshirt1.jpg, 0
tshirt2.jpg, 0
......
ankleboot999.jpg, 9

    def __init__(self, annotations_file, img_dir, transform=None, target_transform=None):
        self.img_labels = pd.read_csv(annotations_file)
        self.img_dir = img_dir
        self.transform = transform
        self.target_transform = target_transform

 

Parameter

• annotations_file: 이미지 파일 이름과 레이블이 포함된 CSV 파일의 경로.
• img_dir: 이미지 파일이 저장된 디렉토리 경로.
• transform: 이미지에 적용할 변형(전처리).
• target_transform: 레이블에 적용할 변형(전처리).
• self.img_labels: CSV 파일을 읽어와 데이터프레임으로 저장.
• self.img_dir: 이미지 디렉토리 경로 저장.
• self.transform: 이미지 변형 저장.
• self.target_transform: 레이블 변형 저장.

 

__len__

• __len__ 함수는 데이터셋의 샘플 개수를 반환합니다.

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

 

• len(self.img_labels): CSV 파일에 저장된 레이블의 개수를 반환합니다.

 

__getitem__

• __getitem__ 함수는 주어진 인덱스 idx 에 해당하는 샘플을 데이터셋에서 불러오고 반환합니다.
• 인덱스를 기반으로, 디스크에서 이미지의 위치를 식별하고, read_image 를 사용하여 이미지를 텐서로 변환하고,self.img_labels 의 csv 데이터로부터 해당하는 정답(label)을 가져오고, (해당하는 경우) 변형(transform) 함수들을 호출한 뒤, 텐서 이미지와 라벨을 Python 사전(dict)형으로 반환합니다.

def __getitem__(self, idx):
    # 이미지 파일의 경로를 생성합니다. self.img_labels의 첫 번째 열에는 파일 이름이 있습니다.
    img_path = os.path.join(self.img_dir, self.img_labels.iloc[idx, 0])
    
    # 이미지를 읽어와 텐서로 변환합니다.
    image = read_image(img_path)
    
    # 레이블을 가져옵니다. self.img_labels의 두 번째 열에는 레이블이 있습니다.
    label = self.img_labels.iloc[idx, 1]
    
    # transform이 정의되어 있다면, 이미지를 변형합니다.
    if self.transform:
        image = self.transform(image)
    
    # target_transform이 정의되어 있다면, 레이블을 변형합니다.
    if self.target_transform:
        label = self.target_transform(label)
    
    # 이미지와 레이블을 포함하는 샘플을 사전(dict) 형식으로 반환합니다.
    sample = {"image": image, "label": label}
    return sample

Parameter

 

• img_path: 주어진 인덱스에 해당하는 이미지 파일의 경로를 생성합니다.
• read_image(img_path): 이미지 파일을 읽어와 텐서로 변환합니다.
• label: 주어진 인덱스에 해당하는 레이블을 가져옵니다.
• if self.transform: 이미지 변형이 지정된 경우, 변형을 적용합니다.
• if self.target_transform: 레이블 변형이 지정된 경우, 변형을 적용합니다.
• sample = {"image": image, "label": label}: 이미지와 레이블을 포함하는 샘플을 사전(dict) 형식으로 반환합니다.

---

DataLoader로 학습용 데이터 준비하기

• Dataset 클래스는 데이터셋의 특징(feature)을 가져오고, 각 샘플에 정답(label)을 지정하는 일을 합니다.
• 모델을 학습할 때는 데이터를 미니배치(minibatch)로 전달하고, 매 에폭(epoch)마다 데이터를 다시 섞어서 과적합(overfit)을 방지하며, 멀티프로세싱을 사용하여 데이터 검색 속도를 높이려 합니다.
• DataLoader는 이러한 복잡한 과정들을 추상화한 간단한 API를 제공하여, 데이터를 효율적으로 불러올 수 있는 순회 가능한 객체(iterable)를 만듭니다.

DataLoader 생성

먼저 torch.utils.data 모듈에서 DataLoader 클래스를 가져와 학습용 데이터와 테스트 데이터를 위한 DataLoader를 생성합니다.

from torch.utils.data import DataLoader

# 학습 데이터용 DataLoader
train_dataloader = DataLoader(training_data, batch_size=64, shuffle=True)

# 테스트 데이터용 DataLoader
test_dataloader = DataLoader(test_data, batch_size=64, shuffle=True)

 

Parameter

• batch_size=64: 한 번에 64개의 샘플을 가져옵니다.
• shuffle=True: 에폭마다 데이터를 섞어줍니다.

 

DataLoader를 통해 데이터 순회하기

• DataLoader 에 데이터셋을 불러온 뒤에는 필요에 따라 데이터셋을 순회(iterate)할 수 있습니다.
• 아래의 각 순회(iteration)는 (각각 batch_size=64 의 특징(feature)과 정답(label)을 포함하는train_features 와 train_labels 의 묶음(batch)을 반환합니다.
• 또한 shuffle=True 로 지정했으므로, 모든 배치를 순회한 뒤 데이터가 섞입니다.

# DataLoader에서 첫 번째 배치 가져오기
train_features, train_labels = next(iter(train_dataloader))

# 배치의 크기 출력
print(f"Feature batch shape: {train_features.size()}")
print(f"Labels batch shape: {train_labels.size()}")

# 첫 번째 이미지와 라벨을 가져와서 시각화
img = train_features[0].permute(1, 2, 0) / 2 + 0.5  # 정규화 해제 및 (C, H, W) -> (H, W, C) 변경
label = train_labels[0]
plt.imshow(img.numpy())
plt.title(f"Label: {label}")
plt.axis("off")
plt.show()

print(f"Label: {label}")

 

 

Feature batch shape: torch.Size([64, 3, 32, 32])
Labels batch shape: torch.Size([64])

Label: 4

Parameter

 

• next(iter(train_dataloader)): DataLoader에서 첫 번째 배치를 가져옵니다.
• train_features.size(): 특징(feature) 배치의 크기를 출력합니다.
• train_labels.size(): 정답(label) 배치의 크기를 출력합니다.
• train_features[0]: 첫 번째 이미지를 가져옵니다.
• permute(1, 2, 0): 이미지의 차원을 (H, W, C) 형식으로 변경하여 시각화합니다.
• img.numpy(): 이미지를 NumPy 배열로 변환하여 plt.imshow로 시각화합니다.
• plt.title(f"Label: {label}"): 이미지의 라벨을 제목으로 표시합니다.
• plt.axis("off"): 축을 끕니다.

---

 

더 자세한 내용을 보고 싶으시면 아래 링크에 들어가서 공식 문서를 참고하세요!

• torch.utils.data 관련 PyTorch 공식 문서

torch.utils.data — PyTorch 2.4 documentation