RuntimeError: Expected target size [4, 256], got [4, 256, 256, 1]

1

Hi everyone, I am aware that this sort of questions are already discussed but I couldnt make it work, so I have 3 channel 256x256 images and 7 classes for a semantic segmentation task.
RuntimeError: Expected target size [4, 256], got [4, 256, 256, 1] I am getting this error write now.

class CustomImageDataset(Dataset):
    def __init__(self, mask, img, transform=None):
        self.mask = mask
        self.img = img
        self.transform = transform

    def __len__(self):
        return np.unique(mask).size

    def __getitem__(self, idx):
        image = self.img[idx]
        mask = torch.from_numpy(self.mask[idx])
        mask = mask.type(torch.LongTensor)

        if self.transform:
            image = self.transform()(image)
 
        return image, mask

Here is my dataset creator code. My inputs to this dataset creator is (num_img, 256,256,3) and (num_masks, 256,256,1) .

Edit : I was also getting “RuntimeError: only batches of spatial targets supported (3D tensors) but got targets of dimension: 4” this error and I handled it with argmax operation (RuntimeError: only batches of spatial targets supported (3D tensors) but got targets of dimension: 4) however it is causing 0 gradient error too. I feel quite lost…

2

The input data shape looks wrong as PyTorch uses the channels-first memory layout by default. The inputs are thus expected to have the shape [batch_size, channels, height, width] and you might need tom .permute the input tensor.

In a multi-class segmentation use case using nn.CrossEntropyLoss the model outputs are expected to have the shape [batch_size, nb_classes, height, width] while the targets should have the shape [batch_size, height, width] containing class indices in the range [0, nb_classes-1].

1 Like