Advanced 2D Convolution - CNNs & Computer Vision

Advanced 2D Convolution

What you'll build

A batched, multi-channel 2D convolution with padding, stride and groups: the computation behind F.conv2d(input, weight, padding=..., stride=..., groups=...) (no bias).

Background

Real conv layers mix many input channels into many output channels, downsample with a stride, and pad to control the output size. Grouped convolution splits the channels into independent groups: AlexNet used 2 groups to fit on two GPUs, ResNeXt uses 32, and with one group per channel you get the depthwise convolution of MobileNet and ConvNeXt, which costs 1/C_in of a full conv.

Layout (NCHW, as in PyTorch)

Tensor	Shape
input	(N, C_in, H, W) = (batch, channels, height, width)
kernel	(C_out, C_in // groups, kH, kW)
returned output	(N, C_out, H_out, W_out)

Formula

With padding p, stride s and G = groups, let C_g = C_in // G and O_g = C_out // G. Output channel o belongs to group g = o // O_g and only sees input channels g*C_g ... (g+1)*C_g - 1:

x_pad = input zero-padded by p on each side of H and W
out[n, o, i, j] = sum over c in [0, C_g), a in [0, kH), b in [0, kW) of
                  x_pad[n, g*C_g + c, i*s + a, j*s + b] * kernel[o, c, a, b]

H_out = (H + 2p - kH) // s + 1
W_out = (W + 2p - kW) // s + 1

This is cross-correlation (the kernel is not flipped), like every deep learning framework.

Inputs

• input: float array (N, C_in, H, W).
• kernel: float array (C_out, C_in // groups, kH, kW); kernels may be non-square.
• padding (int, default 0), stride (int >= 1, default 1), groups (int >= 1, default 1). C_in and C_out are both divisible by groups.

Expected return format

A float NumPy array of shape (N, C_out, H_out, W_out).

Special cases

• groups = 1: standard convolution (every output channel sees every input channel).
• groups = C_in = C_out: depthwise convolution (one kernel per channel).

Constraints

NumPy only (np.pad and np.einsum are fine); no SciPy or deep learning libraries. Loops over the kH * kW kernel taps (and over groups) are fine; avoid looping over every output pixel if you can.

Follow-up

• How many parameters and multiply-adds does this layer have? How does a depthwise separable conv (depthwise 3x3 + pointwise 1x1) compare with a full 3x3?
• How would you implement it with im2col and one batched matmul?

Input: input (1, 4, 5, 5), kernel (8, 2, 3, 3), padding=0, stride=1, groups=2
Output: output shape (1, 8, 3, 3)

Input: input = np.ones((1, 1, 3, 3)), kernel = np.ones((1, 1, 2, 2)), padding=1, stride=2
Output: [[[[1., 2.], [2., 4.]]]]

import numpy as np

def conv2d_advanced(input, kernel, padding=0, stride=1, groups=1):
    """
    2D convolution (cross-correlation) with padding, stride and groups, NCHW layout.

    Args:
        input: Input array, shape (N, C_in, H, W)
        kernel: Kernels, shape (C_out, C_in // groups, kH, kW)
        padding: Zeros added to each side of H and W (default 0)
        stride: Step between kernel positions (default 1)
        groups: Number of channel groups (default 1)

    Returns:
        output: Float array, shape (N, C_out, H_out, W_out) with
                H_out = (H + 2*padding - kH) // stride + 1 (same for W)
    """
    # Your code here
    pass