My gradients of feature map is of the size: torch.Size([1, 64, 128, 128]). My filters for the corresponding convolutional layer is: torch.Size([64, 3, 5, 5]). How do I multiply these two tensors to get the input: torch.Size([1, 3, 128, 128])?
An example implementation of 2D conv layer using F.unfold() is here: codesmart/custom_conv.py at master · InnovArul/codesmart · GitHub
To answer your question, to get grad_input, try following these lines:
# grad_input depends on weights and grad_output
# unfold grad output for grad input calculation
unfold_gradout = F.unfold(grad_output, kernel_size=kernel_size, dilation=1, padding=kernel_size//2, stride=1)
unfold_gradout = rearrange(unfold_gradout, 'b (out_channels kh kw) (h w) -> b out_channels kh kw h w', out_channels=out_channels, kh=kernel_size, kw=kernel_size, h=h, w=w)
weights_for_gradinput = weights.flip(dims=[2,3])
grad_input = torch.einsum('oipq,bopqhw->bihw', weights_for_gradinput, unfold_gradout) # p, q = kernel sizes | h,w = input & output height, width