[PyTorch 에러] Pytorch RuntimeError: stack expects each tensor to be equal size

PyTorch로 이미지 처리를 하다보면 한 번씩 마주해볼 수 있는 에러인데요.

기존에 쉽게 불러올 수 있는 CIFAR-10 이나 MNIST 같은 데이터가 아닌 개인 데이터들을 사용하다보면 각 이미지마다 Size가 다를 수 있습니다. 그래서 torch의 Dataset, DataLoader를 구현해서 사용하고, Dataset 내에서 torchvision이나 albumentation을 통해 Transforms를 정의해서 Resize, Flip 등의 Augmentation을 하실 텐데요.

이 때, 각 이미지마다 Size가 다른 경우에 동일한 Size로 Resize를 시켜주지 않으면, 위와 같은 오류가 발생하게 됩니다. 이는 Array나 Tensor는 각 차원이 모두 동일해야 Batch 형태로 묶어줄 수 있기 때문에 발생합니다.

 

 

class Example_Dataset(nn.Module):
	def __init__(self, df, img_folder_path, transforms=None):
    	self.images = glob.glob('your_folder_path')
        self.label = df['label'].values
        self.transforms = transforms

	def __len__(self):
    	return len(self.images)
        
	def __getitem__(self, idx):
        image = cv2.imread(self.images[idx])
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        label = self.label[idx]

        if self.transforms is not None:
          augmented = self.transforms(image=image)
          image = augmented['image']   
          
        return image, label
        

transforms = albu.Compose([
albu.Resize(img_size, img_size, always_apply=True),
])

dataset = Example_Dataset(df, your_folder_path, transforms)
dataloader = DataLoader(dataset, batch_size=16)

따라서, 위와 같은 예제 코드처럼 image에 대해서 transforms를 적용해주어야 합니다.

* 위 코드의 transforms 부분은 albumentation 기준으로 작성한 것이라 torchvision은 아주 조금 다를 수 있습니다.