ML coding · Deep learning layers
2-D convolution with stride and padding
Given a single-channel image (H × W) and a kernel (kh × kw),
pad the image with padding zeros on every side, then slide the kernel with step
stride in both directions. Each output entry is the sum of the elementwise product of
the kernel with the window under it (no kernel flip). Return the output, of size
⌊(H + 2p − kh)/s⌋ + 1 by ⌊(W + 2p − kw)/s⌋ + 1.
Examples
Input image = [[1, 2, 3], [4, 5, 6], [7, 8, 9]], kernel = [[1, 0], [0, -1]], stride = 1, padding = 0
Output [[-4, -4], [-4, -4]]
Every 2 × 2 window: top-left minus bottom-right, always −4.
Input image = [[1, 2], [3, 4]], kernel = [[1, 1, 1], [1, 1, 1], [1, 1, 1]], stride = 1, padding = 1
Output [[10, 10], [10, 10]]
Padding 1 keeps the size: each output is the sum of a 3 × 3 neighbourhood.
Input image = [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]], kernel = [[1, 1], [1, 1]], stride = 2, padding = 0
Output [[14, 22], [46, 54]]
Stride 2: non-overlapping windows, like sum pooling.
Constraints
- 1 ≤ H, W ≤ 128, 1 ≤ k_h, k_w ≤ H + 2p and W + 2p
- 1 ≤ stride ≤ 4, 0 ≤ padding ≤ 3
- Checked to a relative tolerance of 10⁻⁹
Hints
Hint 1
Work out the output size first: ⌊(H + 2p − k)/s⌋ + 1 in each direction.
Hint 2
np.pad(image, padding) pads with zeros on all four sides. Output (i, j) reads the window starting at row i·s, column j·s.
Hint 3
To remove the loops, sliding_window_view(x, kernel.shape)[::s, ::s] gives every window; einsum('ijkl,kl->ij', windows, kernel) sums each one against the kernel.
Where candidates go wrong
- Flipping the kernel. That is mathematical convolution; deep learning layers do not flip.
- Starting the window for output (i, j) at (i, j) rather than (i·s, j·s) when stride > 1.
- An off-by-one in the output size, which shows up only when (H + 2p − k) is not divisible by s.
- Padding only one side.
The idea
Pad, then for each output position (i, j) take the window x[i·s : i·s + kh, j·s : j·s + kw] and sum its product with the kernel. The loop version is correct and readable.
The vectorised version builds a view of every kh × kw window at once, without copying, then keeps every s-th one in each direction. A single einsum contracts the last two axes against the kernel. Slicing [::s] on the full set of windows gives exactly the ⌊(H + 2p − k)/s⌋ + 1 positions the formula promises.
Time O(H′W′k_hk_w) for an H′ × W′ output · Space O(H′W′) beyond the padded input
Reference solution
import numpy as np
from numpy.lib.stride_tricks import sliding_window_view
def conv2d(image, kernel, stride, padding):
x = np.pad(np.asarray(image, dtype=float), padding) # zeros on all four sides
k = np.asarray(kernel, dtype=float)
windows = sliding_window_view(x, k.shape)[::stride, ::stride] # (out_h, out_w, kh, kw), no copy
return np.einsum('ijkl,kl->ij', windows, k)Follow-ups an interviewer asks
- Multiple input and output channels: input (C, H, W), kernels (F, C, k, k). One einsum still does it.
- Dilation: the kernel taps are spaced d apart. How does the output size change?
- Write the backward pass: the gradient with respect to the input is a transposed convolution.
- im2col: reshape the windows into a matrix and turn the convolution into one matrix product.
Python runs in your browser with NumPy. The first run downloads the runtime, which takes a few seconds.
Why interviewers ask it
Convolution questions check index arithmetic under pressure. The output-size formula, the unflipped kernel and the padding are where answers go wrong, and the follow-up is always to remove the Python loops.
The theory behind it is in the lesson Attention, normalisation and convolution from scratch.
Keep going
Layer normalisation
Normalise each row over its features, then scale by gamma and shift by beta.
Open the problem → Medium15 minBatch normalisation in training mode
Normalise each feature over the batch and update the running mean and variance the way PyTorch does: biased variance to normalise, unbiased to track..
Open the problem → Medium20 minScaled dot-product attention with a causal mask
Implement softmax(QKᵀ/√d)V with an optional causal mask in NumPy.
Open the problem → Medium20 min2-D convolution with stride and padding
Implement a 2-D convolution (cross-correlation) with zero padding and stride, first with loops, then with sliding_window_view and einsum..
Open the problem →