PyTorch 이미지 데이터셋: torchvision.datasets + 현실 데이터 사용법

PyTorch에서 컴퓨터 비전 실험을 할 때 데이터셋 로딩은 첫 번째이자 매우 중요한 단계입니다. PyTorch 생태계에서 torchvision.datasets는 다양한 벤치마크 데이터셋을 손쉽게 불러올 수 있게 해주며, 연구와 개발에서 빠른 프로토타이핑을 가능하게 합니다. 이번 글에서는 torchvision.datasets.CIFAR10을 포함한 주요 데이터셋, 각 데이터셋의 의의, 그리고 사용자 현실 데이터를 사용할 때의 방법까지 자세히 설명하겠습니다.

---

📦 torchvision.datasets란?

torchvision.datasets는 PyTorch에서 제공하는 표준 데이터셋 모듈로, 다양한 공개 이미지 데이터셋을 쉽게 불러와 사용할 수 있게 해줍니다. 이 모듈은 데이터셋 클래스, 다운로드 관리, 전처리 transform 연결까지 한 번에 처리할 수 있어 매우 편리합니다.

데이터셋은 보통 아래처럼 불러옵니다:

import torchvision.datasets as datasets

cifar10 = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)

• root: 데이터셋을 저장할 경로
• train: 학습용(True) / 테스트용(False)
• download: 데이터가 없으면 다운로드 여부
• transform: 전처리(transform) 파이프라인

---

🔥 torchvision에서 제공하는 주요 데이터셋과 의의

🖼 CIFAR10, CIFAR100

• 10/100개의 클래스, 32x32 RGB 이미지
• 의의: 소규모 이미지 분류 벤치마크로, 빠른 실험 및 비교용 표준 데이터셋

CIFAR10

✏️ MNIST, FashionMNIST

• MNIST: 손글씨 숫자 (28x28 회색조)
• FashionMNIST: 의류 이미지 (28x28 회색조)
• 의의: 저해상도, 단순한 데이터로 기본적인 분류 모델 테스트용

MNIST

📸 ImageNet

• 1000개 클래스, 대규모 고해상도 이미지
• 의의: 대규모 학습 및 transfer learning 표준 데이터셋

ImageNet

🐶 STL10

• 10개 클래스, 96x96 컬러 이미지
• 의의: 제한된 라벨 데이터 + unlabeled data로 semi-supervised 실험에 유용

STL10

🚦 COCO, Pascal VOC

• 객체 탐지, 분할, 키포인트 예측용 데이터셋
• 의의: 컴퓨터 비전 고급 과제용, 연구 benchmark로 매우 널리 사용

😃 CelebA

• 얼굴 이미지, 속성 라벨 포함
• 의의: 얼굴 속성 분류, 얼굴 검출, fairness 연구에 활용

---

📥 torchvision 외부에서 현실 데이터 사용하기

현실 데이터(예: 자체 수집 이미지)를 사용할 때는 ImageFolder 또는 Custom Dataset을 사용합니다.

📁 ImageFolder 사용 (폴더 구조 기반)

data/
 ├── class1/
 │    ├── img1.jpg
 │    └── img2.jpg
 └── class2/
      ├── img3.jpg
      └── img4.jpg

from torchvision import datasets, transforms

transform = transforms.Compose([
    transforms.Resize(224),
    transforms.ToTensor()
])

custom_dataset = datasets.ImageFolder(root='./data', transform=transform)

👉 자동으로 폴더명을 클래스 라벨로 인식, 파일을 로딩함

⚙️ Custom Dataset 만들기

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

class MyDataset(Dataset):
    def __init__(self, img_dir, transform=None):
        self.img_dir = img_dir
        self.transform = transform
        self.img_names = os.listdir(img_dir)

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

    def __getitem__(self, idx):
        img_path = os.path.join(self.img_dir, self.img_names[idx])
        image = Image.open(img_path)
        if self.transform:
            image = self.transform(image)
        return image

👉 데이터셋 구조와 라벨링 방식이 복잡할 때는 직접 Dataset 클래스를 정의하면 유연하게 처리 가능

---

✨ 결론

torchvision.datasets는 강력하고 편리한 벤치마크 데이터셋 모듈로, 빠른 실험과 연구 성능 비교에 필수적인 도구입니다. 하지만 현실 데이터에서는 ImageFolder 또는 Custom Dataset 클래스를 사용해 맞춤형 파이프라인을 구축할 수 있습니다. 데이터셋 선택은 실험 목표, 모델 복잡성, 데이터 규모에 따라 결정되므로 항상 목적에 맞게 설계하세요.

추가로 데이터셋별 transform 예제나 Custom Dataset 설계 팁이 필요하다면 언제든 문의해 주세요!