[Pytorch 오류해결] RuntimeError: stack expects each tensor to be equal size, but got

 

발생 에러 메시지

return torch.stack(batch, 0, out=out)
RuntimeError: stack expects each tensor to be equal size, but got [3, 64] at entry 0 and [4, 64] at entry 1

 

데이터로더로부터 데이터를 가져와 학습시키기 전 발생한 오류다. 텐서 사이즈가 동일하지 않다고 한다. 원인은 배치 사이즈 별로 데이터가 묶이지 않아 발생한다고 한다. 직접적으로 텐서 사이즈를 조작 해주어도 잘 해결되지 않았다. 해결 방법은 DataLoader에 collate_fn 파라미터를 사용해줌으로써 학습을 시작할 수 있었다. collate_fn 파라미터는 텐서 간의 사이즈를 조절하는 역할을 한다. 아래와 같이 함수 하나를 만들어 주고 DataLoader(dataset, ..., collate_fn = collate_fn)과 같은 형식으로 입력해주면 문제 해결 가능하다.

 

def collate_fn(batch):
    return tuple(zip(*batch))