DataLoader의 collate_fn, 서로 다른 샘플의 크기를 하나의 배치로 묶는 방법

🤔 Problem

작업을 하다가 DataLoader를 선언하는 코드 부분에서 처음 보는 Argument가 있었습니다. 이 Argument에 대해서 공부를 하면서 '이건 언젠가 유용하게 쓰일 기능이다!'라고 판단이 들어 글을 기록하게 되었습니다.

 

우리는 대부분 모델에 학습을 시킬 때, 각 샘플의 shape이 거의 다 같도록 전처리를 해서 학습을 시키기 때문에 사실 이 기능이 흔하게 쓰이지는 않을 것입니다. 하지만, 특수한 경우에는 각 데이터 샘플의 shape이 똑같이 처리될 수 없는 경우들이 있습니다. 예를 들어, Object Detection에 관한 프로젝트를 한다고 가정하면, 각 이미지에 대해 Bounding Box의 수가 모두 같나요? 거의 대부분 그렇지는 않습니다.

 

이러한 상황에 대해서 하나의 데이터 샘플의 크기는 모두 제각각입니다. 즉, 레이블의 크기가 다양해진다는 말입니다. (물론, 사용하려는 모델에 따라 문제가 되지도 않을 수도 있죠) 그러면, 각 데이터 샘플의 크기가 다른 게 왜 문제가 되나요?

 

배치 단위로 학습을 하기 위해서는 모델에 배치를 넘겨줄 때, 여러 개의 샘플을 묶은 배치를 하나의 텐서로 모델에 넘겨주게 됩니다. 그래서, 하나의 텐서로 묶어주려면 각 샘플의 크기가 완전히 같아야 합니다. (물론, 배치의 사이즈가 하나라면 제각각이어도 됩니다.)

---

😀 Solution

다시 말해, 배치를 가져오는 상황에서 각 샘플의 크기가 다르다면, 제일 크기가 작은 샘플에 대해서 Zero-Padding을 하든, 특정 Constant로 값들을 채워 넣어 크기를 맞추어주든 추가적인 작업을 필요로 합니다. 이러한 작업을 해줄 수 있는 것이 바로 DataLoader의 collate_fn이라는 Argument입니다.

dataloader = DataLoader(dataset, batch_size=batch_size, collate_fn=collate_fn)

위와 같이 선언할 수 있으며, Argument로는 각 텐서를 어떻게 합칠 것인지에 대한 함수를 넣어주어야 합니다. 그러면, 어떻게 함수를 선언해야 하는지 알아봅시다.

 

collate_fn의 Argument로는 batch가 들어오게 됩니다. 이 Batch는 Dataset의 __getitem__ 메서드의 리턴 값들로 리스트를 구성하게 됩니다.

 

아래와 같은 데이터셋 코드가 있다고 합시다.(레퍼런스 링크) 

class CustomDataset(Dataset):
    def __init__(self, num=10):
        self.num = num
    
    def __len__(self):
        return self.num

    def __getitem__(self, idx):
        x = torch.tensor([idx] * (idx+1))
        y = torch.tensor([idx])
        return x, y

레이블은 전부 같지만, input으로 들어오는 텐서의 크기는 다 다르네요.

 

만약, 아래와 같이 Batch Size가 2인 DataLoader를 선언해 줄 것이라면, 아래와 같이 선언할 것입니다.

dl = DataLoader(ds, batch_size=2, collate_fn=collate_fn)

그리고, 이때 input으로 들어갈 텐서의 크기가 제각각이기 때문에 collate_fn을 통해서 '크기가 다른 텐서들을 어떻게 처리할 것이며, 어떻게 하나의 Batch로 합쳐줄 것인가'에 대한 과정을 아래와 같은 함수로 만들어줘야 합니다.

def collate_fn(batch):
    """
    Args:
        - batch : [(sample_1), (sample_2), ...]
        - sample_n : Dataset.__getitem__(n) => Dataset의 getitem 메서드에서 선언한대로 튜플이 들어옴
    Description:
        - 결론적으로 collate_fn의 argument는 Batch에 대한 각 샘플이 한 List에 대한 원소로 구성되어 있으며, 
        - 이 모든 원소(샘플)들을 본 함수 내에서 처리하여, 같은 크기의 Tensor를 갖도록 한 다음에
        - 마지막으로, 이 Tensor들을 모두 합치는 과정을 거치도록 한다.
    """
    
    # Batch 중에 가장 긴 size를 갖는 텐서의 size 찾기
    longest_size = max([sample[0].shape[0] for sample in batch])
    
    # 긴 사이즈의 텐서에 맞춰서 Zero-padding
    x_batch = [F.pad(sample[0], (0, longest_size - sample[0].shape[0]), 'constant', 0) for sample in batch]
    
    # 사이즈를 맞춘 텐서들을 하나의 텐서로 결합
    x_batch = torch.stack([x for x in x_batch], dim=0)
    
    # label 또한 결합
    y_batch = torch.stack([sample[1] for sample in batch], dim=0)
    return x_batch, y_batch

위와 같이 Argument로 들어오는 Batch가 어떤 형태인지를 잘 이해하고 있다면, 원하는 형태의 Batch로 리턴하여 동일한 크기의 텐서로 만들어서 학습을 진행할 수 있습니다.

 

이것을 이해하기 제일 쉽도록 하는 것은 'Batch로 들어오는 Argument는 기존 데이터셋 코드의__getitem__ 메서드가 여러 번 호출된 리턴 값들이 리스트로 담겨있을 뿐이고, 이를 함수 내에서 어떻게 처리하여 원하는 형태로 나타낼 것인가'를 제일 핵심으로 여긴다면 어렵지 않습니다.

📄 원래의 DataLoader로 Batch를 하나씩 출력했을 때

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

📄 collate_fn을 구성하고, Batch Size를 2로 설정하여 DataLoader를 출력했을 때

for x, y in dl:
	print(x, y, end='\n---\n', sep='\n')
    
[Output]
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]])
---

---

📂 Reference

https://seokhee0516.tistory.com/entry/Pytorch-collatefn-%EC%9D%B4%EB%9E%80

Pytorch collate_fn 이란?