Convolution with Unfold

Compute a 2D convolution as one matrix multiplication, the way many fast libraries do.

The trick is called im2col. For every output position, take the patch of the input that the kernel covers and lay it out as a column. Then the convolution at all positions is one matrix product of the flattened kernels with those columns. PyTorch does the patch extraction for you with torch.nn.functional.unfold:

cols = F.unfold(x, kernel_size=(kh, kw), padding=padding, stride=stride)
# x: (N, C, H, W)  ->  cols: (N, C * kh * kw, L)

L is the number of output positions, H_out * W_out, where:

H_out = (H + 2 * padding - kh) // stride + 1
W_out = (W + 2 * padding - kw) // stride + 1

Reshape weight from (O, C, kh, kw) to (O, C * kh * kw), multiply it with cols to get (N, O, L), add the bias for each output channel, and reshape to (N, O, H_out, W_out).

Your result must match F.conv2d, and gradients must flow back to x, weight and bias.

Function signature

def conv2d_unfold(x, weight, bias=None, stride=1, padding=0):

Example

Input

x = torch.arange(16.0).reshape(1, 1, 4, 4)
weight = torch.ones(1, 1, 2, 2)
conv2d_unfold(x, weight, stride=2)

Output

tensor([[[[10., 18.],
          [42., 50.]]]])

Reasoning

With stride 2 the 2x2 kernel covers four separate patches. unfold gives cols of shape (1, 4, 4), one column for each patch. The top left patch is 0, 1, 4, 5, and the all-ones kernel adds it up to 10.

Constraints and notes

• x has shape (N, C, H, W). weight has shape (O, C, kh, kw), and the kernel does not have to be square. bias is None or has shape (O,).
• stride and padding are integers.
• Do not call F.conv2d, nn.Conv2d or any other built-in convolution. The tests check the source.
• weight.reshape(O, -1) @ cols broadcasts over the batch. torch.einsum also works.