Pytorch Dataloader에 대해 알아보자 - 1

pytorch의 기능 중 하나인 dataloader는 학습을 위한 data를 GPU에 feed하기 전 데이터의 변환을 책임지는 모듈로써, tensor의 변환과 batch 관리가 주 역할이라 볼 수 있다. 자세한 내용은 아래의 링크에서 확인할 수 있다.

https://pytorch.org/docs/stable/data.html?highlight=dataloader#torch.utils.data.DataLoader 

torch.utils.data — PyTorch 2.0 documentation

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)

 

각 변수들의 내용은 다음과 같다

- dataset은 말 그대로 우리가 만든 dataset이다.

- batch_size는 batch의 크기를 의미한다.

- shuffle은 순서를 섞을지를 결정한다

- sampler/batch_sampler는 index를 control하는 방법을 정해준다. 이 기능을 사용하기 위해서는 shuffle이 False여야  한다.

- num_workers는 데이터를 불러올 때 사용하는 subprocess의 수이다.

- collate_fn은 map-style 데이터셋에서 sample list를 batch단위로 바꾸기 위해 사용하는 기능으로 함수를 입력값으로 받는다. 주로 zero padding이나 variable size 데이터등 데이터 사이즈를 맞추기 위해 많이 사용한다.

- pin_memory는 tensor를 cuda 메모리에 고정한다

- drop_last는 batch사이즈에 따라 마지막 batch의 크기가 안맞을 수 있기 때문에 제거하는 옵션이다

- timeout은 dataloader를 실행하는 제한시간이다

- worker_init_fn은 어떤 worker를 불러올 지 결정한다.

 

오늘은 이 중 몇가지 핵심기능들을 써보고자 한다. 그 기능들을 설명하기 전에, 먼저 오늘 사용할 dataset을 한번 보자. Pytorch에서 제공하는 Iris(붓꽃) dataset이다.

iris = load_iris()
iris_df = pd.DataFrame(iris['data'], columns=iris['feature_names'])
iris_df['target'] = iris['target']
iris_df

위 코드의 실행 결과

그렇다면, 이 dataset을 이용해 dataloader를 사용해보자. 일단 위의 dataframe을 dataset으로 바꿔줘야 한다.

class IrisDataset(Dataset):
    def __init__(self):
        iris = load_iris()
        self.X=torch.tensor(iris["data"])
        self.y=torch.tensor(iris["target"])
        self.feature_names=iris["feature_names"]

    def __len__(self):
        len_dataset = None
        len_dataset = len(self.y)
        return len_dataset

    def __getitem__(self, idx):
        X, y = None, None
        X = self.X[idx]
        y = self.y[idx]
        return X, y
        
        
dataset_iris = IrisDataset()

 

이렇게 하면 dataloader함수에 넣을 dataset인 dataset_iris를 생성하게 된다. 먼저 아무 기능없이 dataloader를 이용해 dataset을 불러오고, 무슨 값이 있는지 확인하기 위해 iter를 사용해보도록 하겠다.

DataLoader(dataset_iris)

next(iter(DataLoader(dataset_iris)))

>>> [tensor([[5.1000, 3.5000, 1.4000, 0.2000]], dtype=torch.float64), tensor([0])]

출력값을 보면 위의 dataframe의 제일 첫줄 값이 그대로 나온것을 볼 수 있다.

이제부터 핵심 기능들을 사용해보자

 

1. batch_size

batch_size는 batch의 크기를 정해주는 변수이다. 가령 batch_size=4라고 정해주면 한번에 크기가 4인 batch를 생성하게 된다. 값을 할당해주지 않으면 기본값인 1로 결정되며, 위의 코드와 같이 하나의 batch에 하나의 데이터만 들어가게 된다.

next(iter(DataLoader(dataset_iris, batch_size=4)))

>>>
[tensor([[5.1000, 3.5000, 1.4000, 0.2000],
         [4.9000, 3.0000, 1.4000, 0.2000],
         [4.7000, 3.2000, 1.3000, 0.2000],
         [4.6000, 3.1000, 1.5000, 0.2000]], dtype=torch.float64),
 tensor([0, 0, 0, 0])]					#하나의 batch에 4개의 데이터가 들어갔다

 

2. shuffle

shuffle은 DataLoader에서 data를 섞어서 사용할지를 설장한다. 기본값은 False이며, True로 설정하면 데이터의 순서를 섞는다.

next(iter(DataLoader(dataset_iris, shuffle=True, batch_size=4)))

>>>
[tensor([[6.5000, 2.8000, 4.6000, 1.5000],
         [5.7000, 2.8000, 4.1000, 1.3000],
         [6.7000, 3.1000, 4.4000, 1.4000],
         [5.1000, 3.5000, 1.4000, 0.3000]], dtype=torch.float64),
 tensor([1, 1, 1, 0])]
 #실행할때마다 결과가 다르게 나온다. 위에서 shuffle을 설정하지 않았을때와 달리 값이 무작위로 나온다

 

3. sampler/batch_sampler

이 변수는 데이터의 index를 원하는 방식대로 조절하기 위한 변수이다. index를 조절해야 하기 때문에, 이 변수를 사용하기 위해서는 shuffle=False로 설정해야 한다. 직접 함수를 만들수도 있고, 이미 있는 sampler를 사용할 수도 있다. 그종류는 다음과 같다.

• SequentialSampler : 항상 같은 순서
• RandomSampler : 랜덤, replacemetn 여부 선택 가능, 개수 선택 가능
• SubsetRandomSampler : 랜덤 리스트, 위와 두 조건 불가능
• WeigthRandomSampler : 가중치에 따른 확률
• BatchSampler : batch단위로 sampling 가능
• DistributedSampler : 분산처리 (torch.nn.parallel.DistributedDataParallel과 함께 사용

자세한 사용법은 다음의 링크에 자세히 설명되어 있다.

https://towardsdatascience.com/pytorch-basics-sampling-samplers-2a0f29f0bf2a

PyTorch [Basics] — Sampling Samplers

 

4. num_workers

데이터를 불러올 때 사용할 subprocess의 수를 표시한다. 이번 예를 위해 크기가 큰 임의의 dataset을 만들고 num_workers=4로 설정해보자

class RandomDataset(Dataset):
    def __init__(self, tot_len=10, n_features=1):
        self.X = torch.rand((tot_len, n_features))
        self.y = torch.randint(0, 3, size=(tot_len, ))  

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

    def __getitem__(self, idx):
        x = torch.FloatTensor(self.X[idx])
        y = self.y[idx]
        return x, y
        
 dataset_big_random = RandomDataset(tot_len=15000)
 
 %%time
for data, label in DataLoader(dataset_big_random, num_workers=1):
    pass
    
>>>
CPU times: user 12.5 s, sys: 3.84 s, total: 16.3 s
Wall time: 31.9 s
    
%%time
for data, label in DataLoader(dataset_big_random, num_workers=4):
    pass
    
>>>
CPU times: user 13.4 s, sys: 4.39 s, total: 17.8 s
Wall time: 35.2 s

해당 함수는 환경에 따라 영향을 받는다. 나의 경우에는 오히려 시간이 더 걸렸다.

 

5. collate_fn

collate_fn은 map-style의 dataset에서 sample list를 batch단위로 바꾸기 위한 변수이다. zero padding이나 variable size등 데이터 사이즈를 맞추는 등의 작업을 할 수 있다.

위에서 만든 RandomDataset을 이용해서 예시 코드를 써보자

dataset_random = RandomDataset(tot_len=10)


def collate_fn(batch):
    print('Original:\n', batch)
    print('-'*100)
    
    data_list, label_list = [], []
    
    for _data, _label in batch:
        data_list.append(_data)
        label_list.append(_label)
    
    print('Collated:\n', [torch.Tensor(data_list), torch.LongTensor(label_list)])
    print('-'*100)
    
    return torch.Tensor(data_list), torch.LongTensor(label_list)
    
    next(iter(DataLoader(dataset_random, collate_fn=collate_fn, batch_size=4))
    
    
>>>
Original:
 [(tensor([0.3808]), tensor(0)), (tensor([0.6095]), tensor(1)), (tensor([0.6847]), tensor(1)), (tensor([0.3562]), tensor(2))]
----------------------------------------------------------------------------------------------------
Collated:
 [tensor([0.3808, 0.6095, 0.6847, 0.3562]), tensor([0, 1, 1, 2])]
----------------------------------------------------------------------------------------------------
(tensor([0.3808, 0.6095, 0.6847, 0.3562]), tensor([0, 1, 1, 2]))

collate_fn을 통해 feature와 label이 각각 뭉친것을 확인할 수 있다.

 

6. drop_last

batch 단뒤로 데이터를 불러올 때 batch_size에 따라 마지막 batch의 크기가 batch_size값과 다를 수 있기 때문에 마지막 batch를 삭제한다. 기본은 False로 설정되어있으며, drop_last=True로 설정할 경우 마지막 batch를 삭제한다.

for data, label in DataLoader(dataset_random, num_workers=1, batch_size=4):
    print(len(data))	#dataset_random은 위에서 만든 길이가 10인 dataset
    
>>>
4
4
2

for data, label in DataLoader(dataset_random, num_workers=1, batch_size=4, drop_last=True):
    print(len(data))

>>>
4
4		#마지막의 크기가 2인 batch가 사라짐

지금까지 DataLoader에 사용가능한 여러 변수들의 기능을 알아보았다.