PyTorch - RuntimeError: invalid argument 0: Sizes of tensors must match except in dimension 0. Got 224 and 244 in dimension 2
I am trying to do semantic segmentation with two classes.
I have 224x224x3 images and 224x224 binary segmentation masks. I am reshaping the masks to be 224x224x1 (I read somewhere that this is the format that I should pass to the model).
When I try to loop through the train data loader it either runs without errors or I get the following error (randomly):
Traceback (most recent call last):
  File &quot;/Users/nikolaykolibarov/Desktop/GATE/rooftop-recognition/rooftop-edge-segmentation-ml/main.py&quot;, line 43, in &lt;module&gt;
    for i, sample in enumerate(train_loader):
  File &quot;/Users/nikolaykolibarov/opt/anaconda3/envs/dtcc/lib/python3.7/site-packages/torch/utils/data/dataloader.py&quot;, line 346, in __next__
    data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
  File &quot;/Users/nikolaykolibarov/opt/anaconda3/envs/dtcc/lib/python3.7/site-packages/torch/utils/data/_utils/fetch.py&quot;, line 47, in fetch
    return self.collate_fn(data)
  File &quot;/Users/nikolaykolibarov/opt/anaconda3/envs/dtcc/lib/python3.7/site-packages/torch/utils/data/_utils/collate.py&quot;, line 79, in default_collate
    return [default_collate(samples) for samples in transposed]
  File &quot;/Users/nikolaykolibarov/opt/anaconda3/envs/dtcc/lib/python3.7/site-packages/torch/utils/data/_utils/collate.py&quot;, line 79, in &lt;listcomp&gt;
    return [default_collate(samples) for samples in transposed]
  File &quot;/Users/nikolaykolibarov/opt/anaconda3/envs/dtcc/lib/python3.7/site-packages/torch/utils/data/_utils/collate.py&quot;, line 55, in default_collate
    return torch.stack(batch, 0, out=out)
RuntimeError: invalid argument 0: Sizes of tensors must match except in dimension 0. Got 224 and 244 in dimension 2 at /tmp/pip-req-build-9oilk29k/aten/src/TH/generic/THTensor.cpp:689

Most of the issues with the same exception I found online don’t have matching numbers like I do (Got 224 and 224). I don’t understand, this exception should occur when they aren’t equal (if I am not wrong).
When I run it, occasionally I get this error a few times, and then it doesn’t occur a few times.
I am not sure if I messed something up with the data types and shapes or it is something else.
Here is the code:
roof_edges_dataset.py
import os
import cv2 as cv
from torch.utils.data import Dataset
from torchvision.transforms import transforms

from utils import create_binary_mask, get_labelme_shapes, plot_segmentation_dataset


class RoofEdgesDataset(Dataset):
    def __init__(self, im_path, ann_path, transform=None):
        self.im_path = im_path
        self.ann_path = ann_path
        self.transform = transform

        self.im_fn_list = sorted(os.listdir(im_path), key=lambda x: int(x.split('.')[0]))
        self.ann_fn_list = sorted(os.listdir(ann_path), key=lambda x: int(x.split('.')[0]))

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

    def __getitem__(self, index):
        im_path = os.path.join(self.im_path, self.im_fn_list[index])
        im = cv.imread(im_path)

        ann_path = os.path.join(self.ann_path, self.ann_fn_list[index])
        ann = create_binary_mask(im, get_labelme_shapes(ann_path))
        ann = ann.reshape(ann.shape[0], ann.shape[1], 1)
        ann = transforms.ToTensor()(ann)

        if self.transform:
            im = self.transform(im)

        return im, ann

main.py:
import torch
import torchvision.transforms as transforms
from torch.utils.data import DataLoader

from roof_edges_dataset import RoofEdgesDataset

# Device
device = torch.device(&quot;cuda&quot; if torch.cuda.is_available() else &quot;cpu&quot;)

# Hyperparameters
in_im_shape = (3, 224, 224)
num_classes = 2  # Edge / Non-edge
learning_rate = 0.001
batch_size = 4
n_epochs = 10

# Data - 60% Train - 20% Val - 20% Test
transformations = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

dataset = RoofEdgesDataset(im_path='data/images', ann_path='data/annotations', transform=transformations)

train_size = int(0.8 * len(dataset))
test_size = len(dataset) - train_size
train_dataset, test_dataset = torch.utils.data.random_split(dataset, [train_size, test_size])

train_size = int(0.8 * len(train_dataset))
val_size = len(train_dataset) - train_size
train_dataset, val_dataset = torch.utils.data.random_split(train_dataset, [train_size, val_size])

train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(dataset=val_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=True)

# Model

# Loss and Optimizer

# Train
for epoch in range(n_epochs):
    for batch_idx, (image, annotation) in enumerate(train_loader):
        image = image.to(device=device)
        annotation = annotation.to(device=device)
        print(image.shape)
        print(annotation.shape)

    break

# Evaluate

I would be also very thankful if you let me know if something is done in a wrong or stupid way and how to do it better.
Thanks in advance!

It seems like your training images have different sizes, e.g. one is 224x224 while another is 244x244.
Add a crop/resize transformation to all images so they can be concatenated together to form a batch.
