Regarding Conv2d input dimensions and weight dimensions

vision
1

Hey there - I’m currently trying to implement a UNET and have written the model, data loader and train loop. When I try to run the train loop, I get the following error message:

RuntimeError: Expected 4-dimensional input for 4-dimensional weight [64, 1, 3, 3], but got 3-dimensional input of size [1, 572, 572] instead

I’ve searched for solutions here as well as the rest of the searchable internet but can’t wrap my mind around it/can’t find my mistake. Apparently this happens when your model at some point expects a specific dimensionality which it does not receive (initially or from the previous layer).

I’m not too “fluent” when it comes to PyTorch so I’m not sure what information you would need, but I’ll take a shot. This is what my DataLoader for my training data looks like:

train_loader = DataLoader(dataset=train_set, shuffle=False, batch_size=1, num_workers=1, pin_memory=True)

If I follow the Traceback, it leads me to this portion of my model:

class UNET(nn.Module):
    def __init__(self):
        super(UNET, self).__init__()
        self.max_pool_2x2 = nn.MaxPool2d(kernel_size = 2, stride = 2)
        self.downwards_01 = double_conv(1,64) # this is the culprit, the first line of my forward()-function

double_conv is a function since I’ll be using that a lot for my UNET to decode and encode. That function looks like this:

def double_conv(input, output):
    conv = nn.Sequential(
        nn.Conv2d(input, output, kernel_size=3),
        nn.ReLU(inplace=True),
        nn.Conv2d(output, output, padding=0, kernel_size=3, stride=1),
        nn.ReLU(inplace=True),
    )
    return conv

Apparently there is something wrong with my Conv2d. What I understand from the error message is that pytorch expects a tensor with four dimensions while my script is returning [1, 572, 572]. If I understand the manual correctly, Conv2d expects four dimensions as [batchsize, channels, heigth, width]. So for some reason either my batchsize or my channels aren’t being taken care of and I can’t figure it out. Can someone help me figure out my mistake?

2

Conv2d expects 4 dimensions (batch num of samples, num of channels, h, w)
if you want to pass single tensor, use unsqueeze(0) to add extra dimension in front of your tensor.

1 Like