Problem using conv2d - wrong tensor input shape

[Nikronic (Nikan Doosti)]

You misunderstood channels with sizes. In Conv2d, you define input/output channel and kernel size and some arbitrary args like padding, not output size. Output size will be determined using kernel_size, stride, padding, etc.
https://pytorch.org/docs/master/generated/torch.nn.Conv2d.html

An image in PyTorch has thre dimensions [channel, height, width]. So for a RGB image, [3, height, width]. If you want to get a 3 channel image as the result, you need to use a convolution that takes images with same channel size of your input which is 3, and 3 channels as the output, nn.Conv2d(3, 3, kernel_size) where kernel_size is the arbitrary size for filters.

I have explained a little about convs in this thread, I hope it helps. Also, there are alot of great demonstrations in youtube, etc. The only thing you need to be careful about is that PyTorch is channel first which means [channel, height, width] meanwhile numpy is [height, width, channel].

Why add an extra dimension to convolution layer weights?

In the following sample class from Udacity’s PyTorch class, an additional dimension must be added to the incoming kernel weights, and there is no explanation as to why in the course. I’ve highlighted this fact by the multi-line comment in __init__: class Net(nn.Module): """ Network containing a 4 filter convolutional layer and 2x2 maxpool layer. """ def __init__(self, weights): """ weights: the kernel values as a tensor (n_kernels, 1, k_height, k_width). …