Pytorch로 dataset 구성하기

딥러닝 모델을 학습시키기 위해 가장 먼저 필요한 준비물은 "데이터"이다. 주어진 데이터를 효과적으로 활용하기 위해, pytorch에서는 Dataset 클래스를 제공하고 있다. pytorch의 Dataset 클래스를 활용해서 학습을 위한 데이터셋을 어떻게 만들 수 있을지 알아보자.

step1. 클래스 정의

사용할 모듈은 다음과 같다.

import os
import torch
from torch.utils.data import Dataset
from PIL import Image
import torchvision.transforms as transforms

커스텀 데이터셋을 만들기 위해 클래스를 만들어보자.

class MyBaseDataset(Dataset):
    def __init__(self, x_data, y_data):
        self.x_data = x_data
        self.y_data = y_data
        
    def __getitem__(self, index): 
        return self.x_data[index], self.y_data[index]
        
    def __len__(self): 
        return self.x_data.shape[0]

위의 코드를 보면 알 수 있듯이 Custom dataset을 위한 클래스를 정의할 때 Dataset 클래스를 상속받아야 한다. 이때 구현해야할 method는 1) __init__, 2) __getitem__, 3) __len__이다.

•  __init__: 생성자, 모델에 사용할 데이터를 담아두는 등, 어떤 인덱스가 주어졌을 때 반환할 수 있게 만드는 초기 작업을 진행한다.
•  __getitem__: 인덱스가 주어졌을 때 해당되는 데이터를 반환할 수 있도록 하는 메서드이다. numpy 배열이나 텐서 형식으로 반환한다. 보통 입력과 출력을 튜플 형식으로 반환하게 된다.
•  __len__: 학습에 사용할 데이터의 총 개수라고 볼 수 있다. 즉 얼마만큼의 인덱스를 사용할지를 반환하는 메서드이다.

 

step2. 객체 생성

위의 클래스에 간단한 데이터를 넣은 예시를 살펴보자.

x_data = torch.arange(100)
y_data = x_data * x_data
dataset = MyBaseDataset(x_data, y_data)
print("dataset example: ", dataset[0], dataset[1], dataset[2])
print("dataset length:", len(dataset))

x는 0~99까지의 값을 가지는 tensor로 정의하였고, y_data는 x의 제곱 값을 가지도록 하였다. 즉 x = tensor([0., 1., 2., ...99]), y = tensor([0., 1., 4., ...99^2]) 로 세팅하였다.위의 실행결과는 다음과 같다.

dataset example:  (tensor(0), tensor(0)) (tensor(1), tensor(1)) (tensor(2), tensor(4))
dataset length: 100

인덱스로 접근이 잘 되며, 길이도 정상적으로 출력되는 것을 알 수 있다.

 

step3. 인덱스에 접근할 때 데이터 불러오도록 만들기

위와 같이 데이터 크기가 얼마 되지 않는다면, 모든 데이터를 dataset 클래스에 저장하면 되지만, ImageNet처럼 1400만개가 넘는 데이터셋을 메모리에 모두 불러오는 것은 사실상 불가능하다.

그래서 일반적으로 생성자에서는 모든 데이터를 불러오지는 않고, 불러올 이미지의 경로를 저장하는 방식을 사용한다. 그리고 데이터셋에 인덱스로 접근할 때 경로에 존재하는 데이터를 불러와 메모리에 적재하면 이 문제를 해결할 수 있다.

그러면 인덱스에 접근할 때 파일 경로로부터 데이터를 불러오는 Dataset 클래스를 구현해보자.

class DogCatDataset(Dataset):
    def __init__(self, data_dir):
        self.data_dir = data_dir
        self.image_path_list = os.listdir(data_dir)
        self.transform = transforms.ToTensor()

    def __getitem__(self, index):
        image_path = os.path.join(self.data_dir, self.image_path_list[index])

        x_data = Image.open(image_path)
        x_data = self.transform(x_data)
        y_data = 1 if "dog" in self.image_path_list[index] else 0

        return x_data, y_data

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

위의 코드는 개(Dog)와 고양이(Cat)을 구별하기 위한 데이터셋이다. 생성자에서는 인수로 넘어온 data_dir 파일들을 데이터로 사용하기 위해 리스트로 저장하고, 이미지를 텐서로 변환하는 transform을 사용한다.

__getitem__으로 인덱스가 넘어오면, 경로에 해당하는 이미지를 열고(Image.open(image_path)), 이를 transform을 사용해 텐서로 변환해 x_data로 반환한다.

타겟 데이터(y_data)는 파일 이름으로부터 추출한다. 파일 이름이 dog.10.jpt, cat.3.jpg 등으로 구분되어있다고 가정했을 때, 이미지가 개인 경우에는 파일 이름에 "dog" 문자열이 포함되어있을 것이므로 고양이인 경우에는 "dog" 문자열이 포함되어있지 않을 것이다. 따라서 위와 같은 if else 문으로 타겟 데이터(y_data)를 생성하였다. 

결론적으로 위와 같이 데이터를 미리 저장해두는 것이 아니라, 인덱스로 접근할 때 해당되는 데이터를 로드하여 사용하게 되면 메모리에 무리 없이 사용할 수 있다. 

 

step4. 커스텀 데이터셋(Custom Dataset)으로 선형 회귀 구현하기

step1~step3의 과정을 통해 커스텀 데이터셋(Custom Dataset)을 만드는 법을 알아보았다. 그럼 이제 커스텀 데이터셋을 사용하여 선형 회귀를 구현해보자.

import torch
import torch.nn.functional as F
from torch.utils.data import Dataset
from torch.utils.data import DataLoader

class CustomDataset(Dataset): 
  def __init__(self):
    self.x_data = [[73, 80, 75],
                   [93, 88, 93],
                   [89, 91, 90],
                   [96, 98, 100],
                   [73, 66, 70]]
    self.y_data = [[152], [185], [180], [196], [142]]

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

  def __getitem__(self, idx): 
    x = torch.FloatTensor(self.x_data[idx])
    y = torch.FloatTensor(self.y_data[idx])
    return x, y

 

이렇게 정의한 CustomDataset 클래스의 객체를 생성하고, torch에서 제공하는 DataLoader 클래스를 사용해 load하면 된다.

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

shuffle은 datset의 순서를 섞어주는 옵션으로, network가 순서까지도 학습을 해버릴 수 있기 때문에 이를 방지하기 위한 장치라고 생각하면 될 것 같다.

nb_epochs = 20
for epoch in range(nb_epochs + 1):
  for batch_idx, samples in enumerate(dataloader):
    # print(batch_idx)
    # print(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, nb_epochs, batch_idx+1, len(dataloader),
        cost.item()
        ))

위의 코드를 실행시켰을 때 아래와 같이 출력이 되며 학습이 진행된다.

Epoch    0/20 Batch 1/3 Cost: 29410.156250
Epoch    0/20 Batch 2/3 Cost: 7150.685059
Epoch    0/20 Batch 3/3 Cost: 3482.803467
... 중략 ...
Epoch   20/20 Batch 1/3 Cost: 0.350531
Epoch   20/20 Batch 2/3 Cost: 0.653316
Epoch   20/20 Batch 3/3 Cost: 0.010318

그리고 임의의 입력 [73, 80, 75]를 학습된 모델에 입력을 해보자.

new_var =  torch.FloatTensor([[73, 80, 75]]) 
# 입력한 값 [73, 80, 75]에 대해서 예측값 y를 리턴받아서 pred_y에 저장
pred_y = model(new_var) 
print("훈련 후 입력이 73, 80, 75일 때의 예측값 :", pred_y)

위의 실행 결과 아래와 같이 출력된다.

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

dataset에 있던 정답 152와 유사한 값이 출력된다.

---

 

이 포스트는 나의 공부를 위해 아래의 글들을 참고하여 정리한 글이다!! 내 글을 읽고 이해가 안되는 부분은 아래 글들을 참고하면 좋을 것 같다.

https://velog.io/@tjdtnsu/PyTorch-%EA%B8%B0%EC%B4%88-Dataset-%EB%A7%8C%EB%93%A4%EA%B8%B0

PyTorch 기초 - Dataset 만들기

https://wikidocs.net/57165

03-07 커스텀 데이터셋(Custom Dataset)