Incorrect gradient with different stride/kernel sizes

[conscell]

@slvrfn
Thank you for the nice tutorial.
I compared the gradients obtained using your implementation with PyTorch's conv2d. For some values of strides and kernel sizes, the results seem to be incorrect.

import numpy as np
from convolution import Conv2D

in_channels = 3
out_channels = 128
kernel_size = 3
stride = 3                               ### <-- Here, the stride is changed
padding = 1
h_in = 12
w_in = 10

h_out = (h_in + 2 * padding - kernel_size) // stride + 1
w_out = (w_in + 2 * padding - kernel_size) // stride + 1

batch_size = (4, in_channels, h_in, w_in)
dout_size = (4, out_channels, h_out, w_out)

np.random.seed(42)

x = np.random.random(batch_size)  # create data for forward pass
dout = np.random.random(dout_size)  # create random data for backward
print('x: ', x.shape)
print('d_out: ', dout.shape)

conv = Conv2D(in_channels, out_channels, kernel_size, stride, padding)

conv_out = conv.forward(x)
print('conv_out: ', conv_out.shape)

db, dw, dx = conv.backward(dout)
print('db: ', db.shape)
print('dw: ', dw.shape)
print('dx: ', dx.shape)

import torch
xx = torch.tensor(x, requires_grad=True)
ww = torch.tensor(conv.weight, requires_grad=True)
bb = torch.tensor(conv.bias, requires_grad=True)
dd = torch.tensor(dout)

f = torch.nn.functional.conv2d(xx, ww, bb, stride, padding)
f.backward(dd)

assert np.allclose(dx, xx.grad.numpy())

assert np.allclose(dx, xx.grad.numpy())
AssertionError

Additionally, for kernel_size = 2 and stride = 2, the result is also incorrect.

[slvrfn]

Hey thanks for reaching out, I'm glad you enjoyed the tutorial!

I created the tut and this code for a class project a few years ago, and for my purposes it "passed". I currently don't have the bandwidth to address the issue you found, but I may get to it later.

In the mean time, feel free to suggest some code changes, so that someone else can learn!

[conscell]

@slvrfn
I am developing a lightweight PyTorch-like automatic differentiation engine μGrad for educational purposes. Your approach with np.einsum is a gem! It’s significantly faster and more elegant than the complex shape massaging I had been using.

I’ve identified and fixed a few issues:

The stride issue can be addressed by dilating the input with the appropriate number of zeros. Without this adjustment, the results are only accurate for strides 1 and 2.

...
# dilate the input if necessary
if dilate != 0:
    working_input = np.insert(working_input, np.arange(dilate, input.shape[2] * dilate) // dilate, 0, axis=2)
    working_input = np.insert(working_input, np.arange(dilate, input.shape[3] * dilate) // dilate, 0, axis=3)

The following issue arises when passing x.shape as output_size, where some elements of dout_windows might point to invalid memory.

dout_windows = getWindows(dout, x.shape, self.kernel_size, padding=padding, stride=1, dilate=self.stride - 1)

A safer solution is to calculate the output shape within the function itself, especially since the first two dimensions aren’t used there.

out_h, out_w = (working_input.shape[2] - kernel_size) // stride + 1, (working_input.shape[3] - kernel_size) // stride + 1

And the last issue is with padding. For dx computation we unconditionally set padding = self.kernel_size - 1. When the remainder of the division by stride isn’t zero, we need to pad the right and bottom of the output with zeros to match the size of x (padded if necessary). Additionally, if padding was applied during forward computation, we omit it from the final output.

...
dx = np.einsum('bohwkl,oikl->bihw', dout_windows, rot_kern)

dx = np.pad(dx, pad_width=((0, 0), 
                           (0, 0), 
                           (0, max(x.shape[2] + 2 * self.padding - dx.shape[2], 0)),
                           (0, max(x.shape[3] + 2 * self.padding - dx.shape[3], 0))))[..., self.padding:x.shape[2] + self.padding, self.padding:x.shape[3] + self.padding]

There may be a more elegant way to handle these things, so I’m not submitting a PR. I’ve tested the corrected version for both forward and backward passes with the following parameter combinations:

kernels = [1, 2, 3, 4, 5, 6, 7]
strides = [1, 2, 3, 4, 5, 6, 7]
paddings = [0, 1, 2, 3, 4]
hs = [12, 15]
ws = [10, 15]