[Pytorch 구조 학습하기]Dataset과 DataLoader

Pytorch의 Dataset

파이토치는 스레딩을 통한 데이터 병렬화, 데이터 증식 및 배치처리 같은 여러 복잡한 작업을 추상화하는 여러 유틸리티 클래스를 제공한다.

쉽게 말해 raw data를 쉽게 클래스로 만들도록 도와주는 도구라고 보면 된다.

Pytorch의 torch.utils.data에서 Dataset 클래스를 상속해서 만든다.

 

Dataset의 기본 구성 요소

Dataset 클래스를 크게 3가지 메서드로 구성된다.

from torch.utils.data import Dataset

class CustomDataset(Dataset):
    def __init__(self,):
        pass

    def __len__(self):
        pass

    def __getitem__(self, idx):
        pass

__init__

데이터의 위치나 파일명과 같은 초기화 작업을 위해 동작한다. csv이나 xml과 같은 데이터를 이때 불러온다.

여기에 이미지를 처리할 transforms들을 Compose해서 정의한다.

 

__len__

Dataset의 길이(최대 요소 수)를 반환하는 데 사용된다.

 

__getitem__

데이터셋의 idx번째 데이터를 반환하는데 사용된다.

일반적으로 특정 인덱스의 raw data를 가져와서 전처리하고 데이터 증강하는 부분이 여기에서 진행된다.

이 방법은 모든 데이터를 메모리에 로드하지 않고 사용한다는 점에서 효율적이다.

 

 

 

---

Pytorch의 DataLoader

 

DataLoader의 기본 구성 요소

DataLoader는 모델 학습을 위해서 데이터를 미니 배치 단위로 제공해주는 역할을 한다.

DataLoader(dataset, batch_size=1, shuffle=False, sampler=None,
           batch_sampler=None, num_workers=0, collate_fn=None,
           pin_memory=False, drop_last=False, timeout=0,
           worker_init_fn=None)

Pytorch Documentation을 확인해보면 다음처럼 나와있다.

인자를 모두 활용하지는 않고 batch_size나 collate_fn 정도를 주로 사용한다.

 

dataset

사용할 데이터셋 클래스

 

batch_size

배치 사이즈

한 배치당 들어있는 데이터의 개수를 의미한다.

 

shuffle

데이터를 섞을 건지 여부
매 에폭마다 데이터를 새로 섞는다.

 

sampler

데이터의 index를 원하는 방식대로 조정한다.

index를 컨트롤하기 때문에 설정하고 싶다면 shuffle은 False여야 한다.

map-style에서 컨트롤하기 위해 사용하며 __len__과 __iter__를 구현하면 된다.

 

num_workers

데이터를 불러올 때 서브프로세스의 개수

일반적으로 num_workers를 높이면 성능이 좋아지지만 무작정 늘린다고 좋진 않다. 데이터를 불러 CPU와 GPU 사이에 많은 교류가 일어나면 오히려 병목이 생길 수 있다.

 

collate_fn

map-style 데이터셋에서 sample list를 batch 단위로 바꾸기 위해 필요한 기능이다.

zero padding이나 variable size 등 데이터 사이즈를 맞추기 위해 많이 사용한다.

((피처1, 라벨1), (피처2, 라벨2))와 같은 배치 단위 데이터가 ((피처1, 피처2), (라벨1, 라벨2))로 바뀐다.

 

next(iter(DataLoader(dataset_random, collate_fn=collate_fn, batch_size=4)))

Original:
 [(tensor([0.0113]), tensor(1)), (tensor([0.2369]), tensor(0)), (tensor([0.7359]), tensor(1)), (tensor([0.4268]), tensor(2))]
----------------------------------------------------------------------------------------------------
Collated:
 [tensor([0.0113, 0.2369, 0.7359, 0.4268]), tensor([1, 0, 1, 2])]
----------------------------------------------------------------------------------------------------
(tensor([0.0113, 0.2369, 0.7359, 0.4268]), tensor([1, 0, 1, 2]))

collate 이후를 확인해보면 feature는 feature끼리, label은 label끼리 합쳐진 것을 확인할 수 있다.

 

 

혹은 다음처럼 배치 내의 데이터끼리 패딩을 줄 수도 있다.

def my_collate_fn(samples):
    collate_X = []
    collate_y = []
    max_len = max([len(sample['X']) for sample in samples])
    for sample in samples:
        diff = max_len-len(sample['X'])
        if diff > 0:
            zero_pad = torch.zeros(size=(diff,))
            collate_X.append(torch.cat([sample['X'], zero_pad], dim=0))
        else:
            collate_X.append(sample['X'])
    collate_y = [sample['y'] for sample in samples]
    return {'X': torch.stack(collate_X),
             'y': torch.stack(collate_y)}

dataloader_example = torch.utils.data.DataLoader(dataset_example, 
                                                 batch_size=2,
                                                 collate_fn=my_collate_fn)
for d in dataloader_example:
    print(d['X'], d['y'])

tensor([[0., 0.],
        [1., 1.]]) tensor([0., 1.])
tensor([[2., 2., 2., 0.],
        [3., 3., 3., 3.]]) tensor([2., 3.])
tensor([[4., 4., 4., 4., 4., 0.],
        [5., 5., 5., 5., 5., 5.]]) tensor([4., 5.])
tensor([[6., 6., 6., 6., 6., 6., 6., 0.],
        [7., 7., 7., 7., 7., 7., 7., 7.]]) tensor([6., 7.])
tensor([[8., 8., 8., 8., 8., 8., 8., 8., 8., 0.],
        [9., 9., 9., 9., 9., 9., 9., 9., 9., 9.]]) tensor([8., 9.])

 

pin_memory

Tensor를 CUDA 고정 메모리에 할당시킨다. 고정된 메모리에서 데이터를 가져오기 때문에 데이터 전송이 훨씬 빠르다.

하지만 일반적으로 많이 사용하지 않는 argument이다.

 

drop_last

batch size에 따라 마지막 batch의 길이가 달라질 수 있다. 예를 들어, 100개의 데이터를 batch size 3으로 나눈다면 마지막 배치는 1개가 될 것이다.

이렇게 batch의 길이가 다르면 loss를 구하기 귀찮아지고, batch size의 크기에 의존도가 높은 함수를 사용하기 까다로워진다.

이럴 때 마지막 batch를 버리는 argument이다.

 

time_out

DataLoader가 data를 불러오는 제한시간이다.

 

worker_init_fn

어떤 worker를 불러올 것인가를 리스트로 전달한다.

 

 

 

MNIST 데이터로 Dataset과 DataLoader 만들어보기

torchvision에서 제공하는 MNIST Dataset을 활용할 수도 있겠지만,

직접 MNIST raw data를 활용해서 Dataset과 DataLoader를 만들어 보자.

 

 

우선 MNIST 데이터의 경로를 저장해 준다.

from torch.utils.data import Dataset, DataLoader, random_split, SubsetRandomSampler, WeightedRandomSampler
import os


BASE_MNIST_PATH = 'data/MNIST/MNIST/raw'
TRAIN_MNIST_IMAGE_PATH = os.path.join(BASE_MNIST_PATH, 'train-images-idx3-ubyte.gz')
TRAIN_MNIST_LABEL_PATH = os.path.join(BASE_MNIST_PATH, 'train-labels-idx1-ubyte.gz')
TEST_MNIST_IMAGE_PATH = os.path.join(BASE_MNIST_PATH, 't10k-images-idx3-ubyte.gz')
TEST_MNIST_LABEL_PATH = os.path.join(BASE_MNIST_PATH, 't10k-labels-idx1-ubyte.gz')

TRAIN_MNIST_PATH = {
    'image': TRAIN_MNIST_IMAGE_PATH,
    'label': TRAIN_MNIST_LABEL_PATH
}

 

 

다음으로 MNIST raw 데이터를 가져오는 함수를 저장해준다. (이해할 필요는 없다.)

# MNIST RAW 데이터를 가져오는 함수
# https://stackoverflow.com/questions/40427435/extract-images-from-idx3-ubyte-file-or-gzip-via-python
def read_MNIST_images(path):
    with gzip.open(path, 'r') as f:
        # first 4 bytes is a magic number
        magic_number = int.from_bytes(f.read(4), 'big')
        # second 4 bytes is the number of images
        image_count = int.from_bytes(f.read(4), 'big')
        # third 4 bytes is the row count
        row_count = int.from_bytes(f.read(4), 'big')
        # fourth 4 bytes is the column count
        column_count = int.from_bytes(f.read(4), 'big')
        # rest is the image pixel data, each pixel is stored as an unsigned byte
        # pixel values are 0 to 255
        image_data = f.read()
        images = np.frombuffer(image_data, dtype=np.uint8)\
            .reshape((image_count, row_count, column_count))
        return images


def read_MNIST_labels(path):
    with gzip.open(path, 'r') as f:
        # first 4 bytes is a magic number
        magic_number = int.from_bytes(f.read(4), 'big')
        # second 4 bytes is the number of labels
        label_count = int.from_bytes(f.read(4), 'big')
        # rest is the label data, each label is stored as unsigned byte
        # label values are 0 to 9
        label_data = f.read()
        labels = np.frombuffer(label_data, dtype=np.uint8)
        return labels

 

 

이제 MNIST 데이터셋을 저장하고 transform할 MyMNISTDataset을 생성한다.

class MyMNISTDataset(Dataset):
    def __init__(self, path, transform, train=True):
      self.path = path
      self.X = read_MNIST_images(self.path['image'])
      self.transform = transform
      self._repr_indent = 3
      self.train = train

      if self.train == True:
        self.y = read_MNIST_labels(self.path['label'])
        self.classes = self.y

    def __len__(self):
        len_dataset = len(self.X)
        return len_dataset

    def __getitem__(self, idx):
        X = self.X[idx]

        if self.transform:
          X = self.transform(X)

        if self.train:
          y = self.y[idx]
          return torch.tensor(X, dtype=torch.double), torch.tensor(y, dtype=torch.long)
        else:
          return torch.tensor(X, dtype=torch.double)

    def __repr__(self):
        '''
        https://github.com/pytorch/vision/blob/master/torchvision/datasets/vision.py
        '''
        head = "(PyTorch HomeWork) My Custom Dataset : MNIST"
        data_path = self._repr_indent*" " + "Data path: {}".format(self.path['image'])
        label_path = self._repr_indent*" " + "Label path: {}".format(self.path['label'])
        num_data = self._repr_indent*" " + "Number of datapoints: {}".format(self.__len__())
        num_classes = self._repr_indent*" " + "Number of classes: {}".format(len(self.classes))

        return '\n'.join([head,
                          data_path, label_path, 
                          num_data, num_classes])

__init__에서 raw data의 path, transform, train 데이터 여부를 저장한다.

__len__은 데이터의 개수를 반환한다.

__getitem__에서 특정 인덱스에 해당하는 데이터를 저장하고, transform 인자가 있다면 transform한다.

학습 데이터일 경우에만(self.train==True) y를 반환한다.

__repr__은 모델에 대한 정보를 출력해주는 메소드이다.

 

 

다음처럼 데이터셋을 생성할 수 있다.

dataset_train_MyMNIST = MyMNISTDataset(path=TRAIN_MNIST_PATH,
                                       transform=transforms.Compose([
                                           transforms.ToTensor()
                                       ]),
                                       train=True
                                       )

(PyTorch HomeWork) My Custom Dataset : MNIST
   Data path: data/MNIST/MNIST/raw/train-images-idx3-ubyte.gz
   Label path: data/MNIST/MNIST/raw/train-labels-idx1-ubyte.gz
   Number of datapoints: 60000
   Number of classes: 60000

 

 

직접 만든 Dataset을 DataLoader에 다음처럼 적용할 수 있다.

dataloader_train_MNIST = DataLoader(dataset=dataset_train_MyMNIST,
                                    batch_size=16,
                                    shuffle=True,
                                    num_workers=4,
                                    )