RuntimeError: stack expects each tensor to be equal size, but got ~

안녕하세요

pytorch를 사용하여 학습 및 추론을 진행 할 때 RuntimeError: stack expects each tensor to be equal size, but got ~과 같은 에러가 DataLoader에서 발생하는 경우가 종종 있습니다.

해당 에러는 사용하는 데이터셋의 이미지의 사이즈가 서로 달라서입니다.

사이즈가 다르게되면 Array나 Tensor의 각 차원이 동일하지 않기 때문에 batch형태로 묶어줄 수 없기 때문에 발생합니다.

따라서 해결하기 위해서는 torchvision에서 제공하는 transforms.resize()함수를 사용하거나 그 외에 다양한 방법으로 Resize를 진행해주시면 됩니다.

 

import torchvision.transforms as transforms

transform=transforms.Compose([
                        transforms.Resize((32,32)),
                        transforms.ToTensor(),
                        transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))
                    ])

긴 글 읽어주셔서 감사합니다~