PyTorch 데이터 불러오기

개요

 딥러닝을 포함한 머신러닝의 근원은 데이터다. 데이터의 수집, 가공, 사용 방법에 따라 모델 성능이 크게 달라지므로, 데이터를 올바르게 불러오는 것은 가장 중요한 단계 중 하나다. 이번 포스팅에서는 PyTorch에서 데이터를 불러오는 다양한 방법을 단계별로 살펴보고, 커스텀 데이터셋을 구현하는 방법까지 알아본다.

 

파이토치 제공 데이터 사용

라이브러리 불러오기

먼저 PyTorch와 이미지 관련 라이브러리, 데이터 처리를 위한 기타 라이브러리를 불러온다.

import torch                           # 파이토치 기본 라이브러리 
import torchvision                     # 이미지 관련 파이토치 라이브러리
import torchvision.transforms as tr    # 이미지 전처리 기능들을 제공하는 라이브러리
from torch.utils.data import DataLoader, Dataset  # 데이터를 모델에 사용할 수 있도록 정리해 주는 라이브러리
import numpy as np                     # 넘파이 기본 라이브러리

 

전처리 및 데이터셋 구성

tr.Compose에 원하는 전처리 과정을 차례대로 넣는다. 예제에서는 이미지 크기를 16x16으로 변환한 후 텐서 타입으로 변환한다.

# 16x16으로 이미지 크기 변환 후 텐서로 변환
transf = tr.Compose([tr.Resize(16), tr.ToTensor()])

이제 torchvision.datasets에서 제공하는 CIFAR10 데이터를 불러온다. root에는 데이터를 다운로드 받을 경로를, train=True이면 학습 데이터를, train=False이면 테스트 데이터를 불러온다. 미리 선언한 전처리를 사용하기 위해 transform=transf를 작성한다.

trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transf)
testset  = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transf)

 

데이터셋의 각 요소는 일반적으로 (이미지, 라벨) 형태의 튜플이다. 예를 들어, trainset[0][0]은 이미지, trainset[0][1]은 라벨이다.

print(trainset[0][0].size())  # 출력: torch.Size([3, 16, 16])

 

파이토치에서는 이미지 한 장이 (채널 수)x(높이)x(너비) 형태로 표현된다. 데이터를 미니 배치(batch) 형태로 모델에 사용할 수 있도록 DataLoader를 이용해 구성한다.

trainloader = DataLoader(trainset, batch_size=50, shuffle=True)
testloader  = DataLoader(testset, batch_size=50, shuffle=False)

print(len(trainloader))  # CIFAR10 학습 데이터 50,000장 → 배치 크기 50 → 1,000개의 배치

 

iter와 next를 이용해 일부 데이터를 확인할 수 있다.

dataiter = iter(trainloader)
images, labels = next(dataiter)

print(images.size())  # 출력: (배치 크기)x(채널 수)x(높이)x(너비)

 

같은 클래스별 폴더 정리한 경우

데이터가 미리 클래스별로 폴더에 정리되어 있다면, PyTorch의 ImageFolder를 이용해 간단하게 불러올 수 있다. 별도의 라벨링 없이 폴더명에 따라 자동으로 라벨이 지정된다.

 

구글 Colab에서 드라이브 마운트

구글 Colab을 사용하는 경우, 드라이브를 마운트해 데이터를 불러올 수 있다.

from google.colab import drive
drive.mount('/content/gdrive')

# 작업 디렉토리 변경 (필요한 경우)
# %cd /content/gdrive/My\ Drive/deeplearningbrov2/pytorch

 

전처리 및 데이터 불러오기

예제에서는 이미지 크기를 128x128으로 변환한 후 텐서로 변환한다.

transf = tr.Compose([tr.Resize((128, 128)), tr.ToTensor()])

# ImageFolder를 사용해 커스텀 데이터셋 불러오기
trainset = torchvision.datasets.ImageFolder(root='./class', transform=transf)
trainloader = DataLoader(trainset, batch_size=2, shuffle=False)

dataiter = iter(trainloader)
images, labels = next(dataiter)

print(images.size(), labels)

 

정형화되지 않은 커스텀 데이터 불러오기

모든 데이터가 라벨별로 정리된 폴더에 있지 않은 경우도 많다. 예를 들어, 이미지 데이터라도 텍스트, 리스트, 배열 등 다양한 형태로 저장될 수 있다.

 

예제 데이터 생성

32x32 크기의 컬러 이미지와 라벨이 각각 100장 있다고 가정한다. python 복사

# (이미지 수)x(높이)x(너비)x(채널 수)
train_images = np.random.randint(256, size=(100, 32, 32, 3))
# (라벨 수)
train_labels = np.random.randint(2, size=(100, 1))

print(train_images.shape, train_labels.shape)

 

커스텀 데이터셋 구현하기

파이토치의 Dataset 클래스를 상속받아 커스텀 데이터셋을 구현할 수 있다. 아래 템플릿을 참고해 데이터를 불러오도록 한다.

 

from torch.utils.data import Dataset

class TensorData(Dataset):

    def __init__(self, x_data, y_data):
        # 이미지 데이터를 FloatTensor로 변환
        self.x_data = torch.FloatTensor(x_data)
        # (이미지 수)x(높이)x(너비)x(채널 수) → (배치 크기)x(채널 수)x(높이)x(너비)
        self.x_data = self.x_data.permute(0, 3, 1, 2)
        # 라벨 데이터를 LongTensor로 변환
        self.y_data = torch.LongTensor(y_data)
        self.len = self.y_data.shape[0]

    def __getitem__(self, index):
        # 인덱스에 해당하는 데이터 반환
        return self.x_data[index], self.y_data[index]

    def __len__(self):
        # 데이터셋의 길이 반환
        return self.len

파이토치에서는 이미지 데이터가 (배치 크기)x(채널 수)x(높이)x(너비) 형태로 사용되므로, 원래 데이터의 (이미지 수)x(높이)x(너비)x(채널 수) 형태를 permute를 통해 변환해야 한다. 커스텀 데이터셋을 구성한 후, DataLoader를 이용해 미니 배치 형태로 데이터를 준비한다.

 

# TensorData 클래스를 이용해 데이터셋 생성
train_data = TensorData(train_images, train_labels)
train_loader = DataLoader(train_data, batch_size=10, shuffle=True)

dataiter = iter(train_loader)
images, labels = next(dataiter)

print(images.size(), labels.size())

 

참고

[파이토치] 실전 인공지능으로 이어지는 딥러닝 - 기초부터 논문 구현까지 강의 | 딥러닝호형 -