CNNs from scratch in NumPy: convolution as a matrix multiply

You can use nn.Conv2d for years without knowing what it does. That is fine until you need to debug a shape error, implement a custom layer, or explain why your model is slow — at which point the abstraction stops helping.

So: convolution forward and backward in NumPy, the im2col trick that makes it fast, and a small network that actually learns.

Convolution, precisely

A kernel slides over the input; at each position you multiply elementwise and sum.

The top-left output: 1(1) + 2(0) + 4(0) + 5(-1) = 1 - 5 = -4.

Three properties make this the right operation for images.

Local connectivity. Each output depends on a small neighbourhood, which matches how visual structure works — edges and textures are local.

Weight sharing. The same kernel applies everywhere. A 3×3 kernel is 9 parameters regardless of image size, against 9 million for a fully-connected layer on a 1000×1000 image.

Translation equivariance. Shift the input, and the output shifts identically. An edge detector detects edges wherever they are.

Output shape

out = floor((in + 2*padding - kernel) / stride) + 1

def conv_out_size(size, kernel, stride=1, padding=0):
    return (size + 2 * padding - kernel) // stride + 1


conv_out_size(32, 3, 1, 0)   # 30  — shrinks by kernel-1
conv_out_size(32, 3, 1, 1)   # 32  — 'same' padding
conv_out_size(32, 3, 2, 1)   # 16  — stride 2 halves it

Memorise the middle one: padding = (kernel - 1) / 2 with stride 1 preserves the spatial size. For 3×3 that is padding 1; for 5×5, padding 2. Most shape bugs are a forgotten padding.

The naive version

Correct, readable, and far too slow to use:

import numpy as np


def conv_naive(X, W, b, stride=1, pad=0):
    """X: (N, C, H, W_in)   W: (F, C, kh, kw)   b: (F,)"""
    N, C, H, W_in = X.shape
    F, _, kh, kw = W.shape

    Xp = np.pad(X, ((0, 0), (0, 0), (pad, pad), (pad, pad)))
    out_h = (H + 2 * pad - kh) // stride + 1
    out_w = (W_in + 2 * pad - kw) // stride + 1
    out = np.zeros((N, F, out_h, out_w))

    for n in range(N):
        for f in range(F):
            for i in range(out_h):
                for j in range(out_w):
                    hs, ws = i * stride, j * stride
                    patch = Xp[n, :, hs:hs + kh, ws:ws + kw]
                    out[n, f, i, j] = np.sum(patch * W[f]) + b[f]
    return out

Six nested loops in Python. For a batch of 32 on 32×32 images with 64 filters that is roughly 50 million iterations. Minutes per forward pass.

im2col: turn convolution into one matmul

The insight that every real implementation uses: convolution is a matrix multiply in disguise. Extract every patch as a column, and the whole layer becomes one W @ X.

input patches                    reshaped kernels
(C*kh*kw, N*out_h*out_w)         (F, C*kh*kw)

  [p1 p2 p3 ... pn]      x       [k1]        =    (F, N*out_h*out_w)
                                 [k2]
   each column is one             ...             then reshape to
   flattened receptive field     [kF]             (N, F, out_h, out_w)

def im2col(X, kh, kw, stride=1, pad=0):
    """(N, C, H, W) -> (C*kh*kw, N*out_h*out_w)"""
    N, C, H, W = X.shape
    out_h = (H + 2 * pad - kh) // stride + 1
    out_w = (W + 2 * pad - kw) // stride + 1

    Xp = np.pad(X, ((0, 0), (0, 0), (pad, pad), (pad, pad)))
    col = np.zeros((N, C, kh, kw, out_h, out_w))

    # Loop over kernel positions (small) instead of output positions (large)
    for y in range(kh):
        y_max = y + stride * out_h
        for x in range(kw):
            x_max = x + stride * out_w
            col[:, :, y, x, :, :] = Xp[:, :, y:y_max:stride, x:x_max:stride]

    return col.transpose(1, 2, 3, 0, 4, 5).reshape(C * kh * kw, -1)

The trick inside the trick: the loops run over kernel positions (9 iterations for a 3×3) rather than output positions (thousands). Everything else is strided slicing, which NumPy does at C speed.

def conv_forward(X, W, b, stride=1, pad=0):
    N, C, H, W_in = X.shape
    F, _, kh, kw = W.shape
    out_h = (H + 2 * pad - kh) // stride + 1
    out_w = (W_in + 2 * pad - kw) // stride + 1

    col = im2col(X, kh, kw, stride, pad)          # (C*kh*kw, N*out_h*out_w)
    W_col = W.reshape(F, -1)                      # (F, C*kh*kw)

    out = W_col @ col + b.reshape(-1, 1)          # (F, N*out_h*out_w)
    out = out.reshape(F, N, out_h, out_w).transpose(1, 0, 2, 3)

    cache = (X, W, b, stride, pad, col)
    return out, cache

Same numbers, 50–200× faster, because BLAS is extremely good at large matrix multiplies. The cost is memory: im2col duplicates every input pixel kh*kw times. That is the trade every framework makes, and it is why cuDNN picks between im2col, FFT and Winograd depending on the shapes.

The backward pass

Three gradients are needed: with respect to the weights, the bias, and the input.

def conv_backward(dout, cache):
    X, W, b, stride, pad, col = cache
    N, C, H, W_in = X.shape
    F, _, kh, kw = W.shape

    # (N, F, out_h, out_w) -> (F, N*out_h*out_w), matching the forward layout
    dout_r = dout.transpose(1, 0, 2, 3).reshape(F, -1)

    db = dout_r.sum(axis=1)                       # (F,)
    dW = (dout_r @ col.T).reshape(W.shape)        # (F, C, kh, kw)

    dcol = W.reshape(F, -1).T @ dout_r            # (C*kh*kw, N*out_h*out_w)
    dX = col2im(dcol, X.shape, kh, kw, stride, pad)

    return dX, dW, db


def col2im(dcol, x_shape, kh, kw, stride=1, pad=0):
    """Inverse of im2col — scatter-ADD, because patches overlap."""
    N, C, H, W = x_shape
    out_h = (H + 2 * pad - kh) // stride + 1
    out_w = (W + 2 * pad - kw) // stride + 1

    col = dcol.reshape(C, kh, kw, N, out_h, out_w).transpose(3, 0, 1, 2, 4, 5)
    dX = np.zeros((N, C, H + 2 * pad, W + 2 * pad))

    for y in range(kh):
        y_max = y + stride * out_h
        for x in range(kw):
            x_max = x + stride * out_w
            dX[:, :, y:y_max:stride, x:x_max:stride] += col[:, :, y, x, :, :]

    return dX[:, :, pad:pad + H, pad:pad + W]

The += in col2im is the part people get wrong. Overlapping patches mean one input pixel contributed to several outputs, so its gradient is the sum of all those contributions. Assignment instead of addition silently produces wrong gradients that still look plausible.

The structural insight: dW = dout @ col.T is a convolution of the input with the output gradient, and dcol = W.T @ dout is a full convolution of the output gradient with the flipped kernel. The backward pass of a convolution is itself convolution — which is why the same optimised kernels serve both directions.

Max pooling

def maxpool_forward(X, size=2, stride=2):
    N, C, H, W = X.shape
    out_h, out_w = (H - size) // stride + 1, (W - size) // stride + 1

    # Treat each channel as its own image, then reuse im2col
    col = im2col(X.reshape(N * C, 1, H, W), size, size, stride, 0)
    col = col.reshape(size * size, -1)

    argmax = col.argmax(axis=0)
    out = col[argmax, np.arange(col.shape[1])]
    out = out.reshape(N, C, out_h, out_w)

    return out, (X, argmax, size, stride)


def maxpool_backward(dout, cache):
    """Gradient flows only to the position that won the max."""
    X, argmax, size, stride = cache
    N, C, H, W = X.shape

    dcol = np.zeros((size * size, dout.size))
    dcol[argmax, np.arange(dout.size)] = dout.transpose(0, 1, 2, 3).ravel()
    dcol = dcol.reshape(-1, dout.size // (N * C) * (N * C))

    return col2im(dcol, (N * C, 1, H, W), size, size, stride, 0).reshape(X.shape)

Max pooling routes the gradient exclusively to the winning position. The losers get exactly zero — pooling is a hard, non-differentiable-at-ties selection, and the argmax recorded in the forward pass is what makes the backward pass possible.

Assembling a network that trains

class ConvNet:
    """conv -> relu -> pool -> conv -> relu -> pool -> fc -> softmax"""

    def __init__(self, in_ch=1, n_classes=10):
        # He initialisation: std = sqrt(2 / fan_in), the correct scale for ReLU
        self.W1 = np.random.randn(16, in_ch, 3, 3) * np.sqrt(2 / (in_ch * 9))
        self.b1 = np.zeros(16)
        self.W2 = np.random.randn(32, 16, 3, 3) * np.sqrt(2 / (16 * 9))
        self.b2 = np.zeros(32)
        self.W3 = np.random.randn(32 * 7 * 7, n_classes) * np.sqrt(2 / (32 * 49))
        self.b3 = np.zeros(n_classes)

    def forward(self, X):
        c = {}
        h, c["c1"] = conv_forward(X, self.W1, self.b1, pad=1)
        h, c["r1"] = relu_forward(h)
        h, c["p1"] = maxpool_forward(h)

        h, c["c2"] = conv_forward(h, self.W2, self.b2, pad=1)
        h, c["r2"] = relu_forward(h)
        h, c["p2"] = maxpool_forward(h)

        c["shape"] = h.shape
        h = h.reshape(h.shape[0], -1)
        return h @ self.W3 + self.b3, c

    def loss(self, X, y):
        logits, c = self.forward(X)

        shifted = logits - logits.max(axis=1, keepdims=True)   # stability
        probs = np.exp(shifted) / np.exp(shifted).sum(axis=1, keepdims=True)
        N = X.shape[0]
        loss = -np.log(probs[np.arange(N), y] + 1e-12).mean()

        dlogits = probs.copy()
        dlogits[np.arange(N), y] -= 1
        dlogits /= N
        return loss, dlogits, c

He initialisation matters more than it looks. ReLU zeroes half its inputs, halving the variance at every layer. sqrt(2/fan_in) compensates for exactly that factor of two. Use Xavier (sqrt(1/fan_in)) with ReLU and activations shrink layer by layer until the deepest layers see nothing — a silent failure that looks like "the model just does not learn".

Gradient checking, which you should always do

Analytical gradients are easy to get subtly wrong. Numerical gradients are slow but unambiguous.

def gradient_check(f, x, analytic_grad, n_checks=10, h=1e-5):
    """f returns a scalar loss. Compare a few random entries."""
    for _ in range(n_checks):
        idx = tuple(np.random.randint(s) for s in x.shape)

        old = x[idx]
        x[idx] = old + h
        fx_plus = f()
        x[idx] = old - h
        fx_minus = f()
        x[idx] = old

        numeric = (fx_plus - fx_minus) / (2 * h)
        analytic = analytic_grad[idx]
        rel_err = abs(numeric - analytic) / max(1e-8, abs(numeric) + abs(analytic))

        status = "ok" if rel_err < 1e-5 else "FAIL"
        print(f"{idx} numeric={numeric:+.6f} analytic={analytic:+.6f} "
              f"rel_err={rel_err:.2e} {status}")

Relative error below 1e-7 is correct. Around 1e-4 there is a real bug. Use the central difference — (f(x+h) - f(x-h)) / 2h — not the forward difference; it is second-order accurate and the extra evaluation is worth it.

The most common bug this catches is exactly the col2im assignment-versus-addition mistake above.

Why this is worth doing once

The framework version is three lines and you should use it in production. But having written the backward pass by hand, several things stop being mysterious: why padding of (k-1)/2 preserves size, why im2col makes convolution memory-hungry, why max pooling gradients are sparse, why He initialisation exists, and why a shape error at layer 4 is almost always a stride you forgot at layer 2.

Frequently asked questions

What is im2col and why does every framework use it?

It rearranges every convolution patch into a column of a matrix, turning the whole layer into a single matrix multiply. BLAS libraries are heavily optimised for large matmuls, so this is 50–200× faster than looping. The cost is memory — each input pixel is duplicated kernel_h × kernel_w times.

Why does my convolution output shrink?

Without padding, a k×k kernel removes k-1 pixels from each spatial dimension. Use padding = (k-1)/2 with stride 1 to preserve the size — padding 1 for 3×3, padding 2 for 5×5.

How does the gradient flow through max pooling?

Only to the position that produced the maximum. Every other position in the window receives exactly zero. That is why the forward pass records the argmax — the backward pass needs it to know where to route the gradient.

Why He initialisation instead of Xavier for CNNs?

ReLU zeroes roughly half of its inputs, which halves the activation variance at every layer. He initialisation uses sqrt(2/fan_in) to compensate for that factor of two. With Xavier, activations shrink layer by layer until deep layers receive almost no signal.

My gradients are wrong but the loss still goes down. How do I find the bug?

Gradient checking with central differences on a handful of random parameter entries. A relative error above about 1e-4 means a real bug. The classic culprit in convolution is using assignment instead of addition when scattering overlapping patches back in col2im.

Next reading on Sythra Articles

• Transformer from scratch in PyTorch
• The attention mechanism, worked out by hand
• Matrix multiplication in NumPy: @ vs *