[실전 예제/이미지 분류/PyTorch] 이미지 분류 튜토리얼: CIFAR-10과 ImageNet으로 PyTorch 데이터셋 만들기

이미지 분류(Image Classification)란?

  이미지 분류(Image Classification)는 딥러닝에서 가장 기본적이면서도 실전 활용도가 높은 컴퓨터 비전 과제입니다. 고양이와 개를 구분하거나, X-ray 사진에서 질병을 탐지하는 문제도 모두 이미지 분류 문제에 해당합니다. 이번 시간에는 PyTorch를 이용하여 이미지 분류 데이터셋을 만드는 방법에 대해 알아보도록 하겠습니다.

 

PyTorch로 CIFAR-10 이미지 분류 데이터셋 만들기

  CIFAR-10은 딥러닝 입문자들이 가장 많이 사용하는 공개 이미지 분류 데이터셋입니다.

CIFAR-10 특징

• 10개 클래스 (비행기, 자동차, 개구리 등)
• 32×32 크기 컬러 이미지
• PyTorch에서 `torchvision.datasets.CIFAR10`로 바로 사용 가능

Pytorch 코드 예제

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

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

 

PyTorch로 ImageNet 또는 커스텀 이미지 분 데이터셋 만들기

  ImageNet은 실제 서비스 수준의 고해상도 이미지 분류를 위해 널리 사용되는 대규모 데이터셋입니다. 모든 사전학습 모델의 기초가 됩니다.

ImageNet 특징

• 1000개 클래스
• 224×224 이상의 컬러 이미지

폴더 구조 (ImageFolder 형식)

data/
  train/
    cat/
    dog/
  val/
    cat/
    dog/

 

PyTorch 코드 예제

from torchvision.datasets import ImageFolder
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

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

train_dataset = ImageFolder('data/train', transform=transform)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

# ImageNet은 ImageFolder 형식이 아니더라도 CIFAR-10과 마찬가지로 PyTorch에서 제공된다.
train_dataset = datasets.ImageNet(root='./data', transform=transform)

 

마무리

  PyTorch를 이용하여 이미지 분류 데이터셋을 어떻게 만드는지 살펴보았습니다. 다음 시간에는 모델 구성 및 학습 방법을 PyTorch로 작성하는 방법을 알아보도록 하겠습니다.

 

관련 내용

• 준비중