[Pytorch] 데이터 로드하기 - Dataset, DataLoader 정리

딥러닝을 포함한 머신러닝의 근원은 데이터이다. 데이터의 수집, 가공, 사용 방법에 따라 모델 성능이 크게 달라질 수 있으며 데이터의 형태는 매우 다양하기 때문에 데이터를 잘 불러오는 것은 가장 중요한 단계 중 하나이다.

Pytorch에서는 데이터를 좀 더 쉽게 다룰 수 있도록 데이터셋(Dataset)과 데이터로더(DataLoader)를 제공한다. 이를 사용하면 미니 배치 학습, 데이터 셔플(shuffle), 병렬 처리까지 간단히 수행할 수 있다.

1. TensorDataset 사용하여 데이터 로드하기

우선 텐서를 입력받아 Dataset의 형태로 변환해주는 TensorDataset을 사용해볼 것이다.

TensorDataset은 기본적으로 텐서를 입력으로 받는다.

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import TensorDataset, DataLoader

x_train  =  torch.FloatTensor([[73,  80,  75], 
                               [93,  88,  93], 
                               [89,  91,  90], 
                               [96,  98,  100],   
                               [73,  66,  70]])  
y_train  =  torch.FloatTensor([[152],  [185],  [180],  [196],  [142]])

dataset = TensorDataset(x_train, y_train)

Pytorch의 Dataset을 만들었다면 DataLoader를 사용할 수 있다. DataLoader는 기본적으로 (데이터셋, 미니 배치의 크기)로 2개의 인자를 입력받는다. 추가적으로, shuffle=True를 인자로 주면 Epoch마다 Dataset을 섞어서 데이터가 학습되는 순서를 바꾼다.

dataloader = DataLoader(dataset, batch_size=2, shuffle=True)

이제 모델과 옵티마이저를 설정 후 훈련을 진행해본다.

model = nn.Linear(3,1)
optimizer = torch.optim.SGD(model.parameters(), lr=1e-5)
n_epochs = 20

for epoch in range(n_epochs + 1):
  for batch_idx, samples in enumerate(dataloader):
    print(f"batch_idx: {batch_idx}")
    print(f"samples: {samples}")
    x_train, y_train = samples
    # H(x) 계산
    prediction = model(x_train)

    # cost 계산
    cost = F.mse_loss(prediction, y_train)

    # cost로 H(x) 계산
    optimizer.zero_grad()
    cost.backward()
    optimizer.step()

    print('Epoch {:4d}/{} Batch {}/{} Cost: {:.6f}'.format(
        epoch, n_epochs, batch_idx+1, len(dataloader),
        cost.item()
        ))

결과는 다음과 같이 출력된다. Cost의 값이 점차 작아지는 것을 확인할 수 있다.

batch_idx: 0
samples: [tensor([[ 96.,  98., 100.],
        [ 93.,  88.,  93.]]), tensor([[196.],
        [185.]])]
Epoch    0/20 Batch 1/3 Cost: 37637.710938
batch_idx: 1
samples: [tensor([[73., 66., 70.],
        [73., 80., 75.]]), tensor([[142.],
        [152.]])]
Epoch    0/20 Batch 2/3 Cost: 4907.253906
batch_idx: 2
samples: [tensor([[89., 91., 90.]]), tensor([[180.]])]
Epoch    0/20 Batch 3/3 Cost: 3287.156006

... 생략 ...

batch_idx: 0
samples: [tensor([[ 73.,  66.,  70.],
        [ 96.,  98., 100.]]), tensor([[142.],
        [196.]])]
Epoch   20/20 Batch 1/3 Cost: 2.096911
batch_idx: 1
samples: [tensor([[93., 88., 93.],
        [89., 91., 90.]]), tensor([[185.],
        [180.]])]
Epoch   20/20 Batch 2/3 Cost: 0.457153
batch_idx: 2
samples: [tensor([[73., 80., 75.]]), tensor([[152.]])]
Epoch   20/20 Batch 3/3 Cost: 0.080784

이제 모델의 입력으로 임의의 값을 넣어 예측값을 확인한다.

new_var =  torch.FloatTensor([[73, 80, 75]])      # 임의의 입력 선언
pred_y = model(new_var)					# 예측값 저장

print("훈련 후 입력이 73, 80, 75일 때의 예측값 :", pred_y) 

>> Result:
훈련 후 입력이 73, 80, 75일 때의 예측값 : tensor([[154.3850]], grad_fn=<AddmmBackward>)

 

2. Pytorch 제공 Dataset 사용하여 데이터 로드하기

import torchvision # 이미지 관련된 파이토치 라이브러리
import torchvision.transforms as tr # 이미지 전처리 기능들을 제공하는 라이브러리
import numpy as np # 넘파이 기본 라이브러리

 

이미지를 16x16 크기로 변환(Resize)한 후 텐서 타입으로 변환(ToTensor)하는 전처리 과정을 선언한다.

tr.Compose를 이용하여 이미지를 전처리하는 과정을 한번에 묶을 수 있다.

tr.Resize는 입력 이미지를 지정한 크기로 변환해주며, tr.ToTensor는 PIL 이미지 또는 numpy 배열을 텐서로 변환해준다.

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

torchvision.datasets에서 제공하는 CIFAR10 데이터를 불러온다.

root에는 데이터를 다운받을 경로를 입력한다.

train=True이면 학습 데이터를 불러오고, train=False이면 테스트 데이터를 불러온다.

전처리 과정 transform의 경우 위에서 미리 선언한 transf를 사용하기 위해 transform=transf로 작성한다.

# root: 데이터를 저장할 경로
# train: 학습 데이터를 불러올지 여부
# download: 다운로드를 할지 여부
# transform: 이미지 변환 과정(Resize, ToTensor)
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transf)
testset = torchvision.datasets.CIFAR10(root'./data', train=False, download=True, transform=transf)

일반적으로 데이터셋은 이미지와 라벨이 동시에 들어있는 튜플(tuple) 형태이다. (이미지, 라벨)

trainset[0]은 학습 데이터의 첫 번째 데이터로 이미지 한 장과 라벨 숫자 하나가 저장되어 있다. 즉, trainset[0][0]은 이미지이며 trainset[0][1]은 라벨이다.

 

현재 이미지 사이즈는 다음과 같이 3x16x16이다.

일반적인 컬러 사진은 RGB 이미지이기 때문에 채널이 3개이고 (높이)x(너비)x(채널 수)로 크기가 표현된다.

하지만 Pytorch에서는 이미지 한 장이 (채널 수)x(높이)x(너비)으로 표현된다.

print(trainset[0][0].size())

>> Result:
torch.Size([3, 16, 16])          # (채널 수)x(높이)x(너비)

 

DataLoader를 사용하여 데이터를 미니배치 형태로 만들어준다. 미니 배치의 크기는 50, shuffle=True로 설정해주었다. CIFAR10의 학습 이미지는 50000장이고 배치 사이즈를 50으로 설정하였으므로 배치의 개수는 1000이다.

이터레이터 객체로 만들어 일부 데이터를 확인해보면 4차원 형태의 학습 데이터를 확인할 수 있다.

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

len(trainloader)     # 1000
train_iter = iter(trainloader)
images, labels = next(train_iter)
print(images.size())	# torch.Size([50, 3, 16, 16]) -> (배치 크기)x(채널 수)x(높이)x(너비)

 

클래스 별로 폴더가 따로 존재할 경우

데이터가 클래스 별로 폴더가 정리되어있는 경우 ImageFolder를 통해 별도의 라벨링 없이 폴더 별로 자동 라벨링이 가능하다. 예를 들어, animal 폴더 내부에 dog, cat, mouse 폴더가 있을 때 ImageFolder의 root에 상위 폴더 ./animal을 입력하면 이미지와 라벨이 정리되어 데이터를 불러온다.

transf = tr.Compose([tr.Resize((128, 128)),tr.ToTensor()]) # 128x128 이미지 크기 변환 후 텐서로 만든다.
trainset = torchvision.datasets.ImageFolder(root='./animal', transform=transf) # 커스텀 데이터 불러온다.
trainloader = DataLoader(trainset, batch_size=2, shuffle=False) # 데이터를 미니 배치 형태로 만들어 준다.

 

커스텀 데이터셋 만들고 로드하기

클래스 별로 폴더가 잘 정리되어있는 경우도 있지만, 실제로는 예외적인 상황이 많다. 함부로 수정할 수 없는 공유 데이터이거나, 이미지 데이터이더라도 이미지가 아닌 텍스트, 배열, 리스트 등으로 저장되어있을 수도 있다. 따라서 커스텀 데이터셋으로 데이터를 로드하는 방법을 알아볼 것이다.

 

커스텀 Dataset 클래스에서는 반드시 __init__,  __len__, __getitem__ 3개의 함수를 구현해야한다.

예시로 32x32 컬러 이미지와 라벨이 각각 100장 있다고 가정한다.

__init__의 data.permute(0, 3, 1, 2)의 경우 pytorch에서는 (배치 크기)x(채널 수)x(너비)x(높이) 데이터가 사용 되므로 기존 데이터 (이미지 수)x(높이)x(너비)x(채널 수)의 형식을 변경하기 위해 진행된 것이다.

import torch
import torchvision.transforms as tr
from torch.utils.data import DataLoader, Dataset
import numpy as np

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

class CustomDataset(Dataset):
    def __init__(self, datas, labels, transform=None):
    	self.datas = datas      # numpy 배열
        self.labels = labels		# numpy 배열
       	self.transform = transform
        self.len = len(labels)
        
    def __getitem__(self, index):
        sample = self.datas[index], self.labels[index] # 특정 인덱스의 샘플 데이터와 라벨 저장
        if self.transform:
        	sample = self.transform(sample)    # transform을 통해 numpy 배열에서 tensor로 변환 및 전처리 작업 진행
        return sample
    
    def __len__(self):
        return self.len # 클래스 내의 들어 온 데이터 개수

이제 전처리 작업을 정의해준다. __call__ 내에 원하는 전처리 작업을 정의한다.

class CustomTransform:
    def __call__(self, sample):
        datas, labels = sample
        datas = torch.FloatTensor(datas)		# 텐서로 변환
        datas = datas.permute(2,0,1)			# shape 변환 : 기존 데이터는 (높이)x(너비)x(채널 수) -> Pytorch에서는 (채널 수)x(높이)x(너비) 데이터가 사용
        labels = torch.FloatTensor(labels)		# 텐서로 변환

        transf = tr.Compose([tr.ToPILImage(), tr.Resize(128),tr.ToTensor(),tr.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
        trans_data = transf(datas)      
        
        return trans_data, labels

CustomDataset과 CustomTransform 클래스를 사용하여 데이터를 불러와 로드한다.

dataset = CustomDataset(images,labels, transform=CustomTransform())
dataloader = DataLoader(dataset, batch_size=15, shuffle=True)

배치 사이즈 만큼 (3, 128, 128) 크기의 이미지를 잘 가져오는 것을 확인할 수 있다.

test_iter = iter(dataloader)
test_sample, test_labels = next(test_iter)
print(test_sample.size()) 			# torch.Size([15, 3, 128, 128])

 

 

Custom Dataset 자세한 내용 참고: https://tutorials.pytorch.kr/beginner/basics/data_tutorial.html

Dataset과 DataLoader

 

 

추가 전처리 방법 참고

https://pytorch.org/vision/stable/transforms.html

Transforming and augmenting images — Torchvision 0.15 documentation