Lesson 33: Convolutional Neural Networks

The layers in Lessons 30-32 treat every input as a flat vector: each pixel gets its own independent weight, with no notion that pixel $(5,5)$ is near pixel $(5,6)$. That throws away the entire spatial structure images have, and it means a fully-connected layer must independently re-learn what an edge looks like at every single position. Convolution — already built from scratch in Lesson 10 — fixes both problems at once.

In [1]:
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt

A conv layer is a projection, restricted and shared

Recall Lesson 29's central idea: $y = w^\top x + b$, a projection. A convolutional layer computes exactly this, but for $x$, it only ever looks at a small local patch (say $3\times3$) instead of the whole image — and, crucially, it reuses the same $w$ at every patch location. Two consequences:

  • Vastly fewer parameters. A $3\times3$ filter has 9 weights, no matter how big the image is — a fully-connected layer over a $16\times16$ image needs 256 weights per output unit.
  • Translation equivariance. Because the same filter slides everywhere, a feature learned at one position is automatically detected at every other position too, for free.

Conv2d, forward and backward, from scratch

This is the same sliding-window computation as Lesson 10's convolution, just with the kernel now a set of learnable weights instead of a fixed, hand-designed one — which means it needs a backward pass too.

In [2]:
def conv2d_forward(img, kernel, bias):
    N, C, H, W = img.shape
    OC, IC, KH, KW = kernel.shape
    OH, OW = H - KH + 1, W - KW + 1
    out = np.zeros((N, OC, OH, OW))
    for n in range(N):
        for oc in range(OC):
            for i in range(OH):
                for j in range(OW):
                    patch = img[n, :, i:i + KH, j:j + KW]
                    out[n, oc, i, j] = np.sum(patch * kernel[oc]) + bias[oc]
    return out

def conv2d_backward(img, kernel, grad_out):
    N, C, H, W = img.shape
    OC, IC, KH, KW = kernel.shape
    OH, OW = grad_out.shape[2], grad_out.shape[3]
    grad_kernel = np.zeros_like(kernel)
    grad_img = np.zeros_like(img)
    grad_bias = grad_out.sum(axis=(0, 2, 3))
    for n in range(N):
        for oc in range(OC):
            for i in range(OH):
                for j in range(OW):
                    patch = img[n, :, i:i + KH, j:j + KW]
                    grad_kernel[oc] += grad_out[n, oc, i, j] * patch
                    grad_img[n, :, i:i + KH, j:j + KW] += grad_out[n, oc, i, j] * kernel[oc]
    return grad_img, grad_kernel, grad_bias

Sanity check against PyTorch

In [3]:
rng = np.random.default_rng(0)
img = rng.normal(size=(1, 1, 8, 8))
kernel = rng.normal(size=(1, 1, 3, 3)) * 0.5
bias = np.array([0.1])

out = conv2d_forward(img, kernel, bias)

img_t = torch.tensor(img, requires_grad=True)
kernel_t = torch.tensor(kernel, requires_grad=True)
bias_t = torch.tensor(bias, requires_grad=True)
out_t = F.conv2d(img_t, kernel_t, bias_t)
print(f'forward max diff: {np.abs(out - out_t.detach().numpy()).max():.2e}')

grad_out = rng.normal(size=out.shape)
grad_img, grad_kernel, grad_bias = conv2d_backward(img, kernel, grad_out)
out_t.backward(torch.tensor(grad_out))

print(f'grad_img max diff:    {np.abs(grad_img - img_t.grad.numpy()).max():.2e}')
print(f'grad_kernel max diff: {np.abs(grad_kernel - kernel_t.grad.numpy()).max():.2e}')
print(f'grad_bias max diff:   {np.abs(grad_bias - bias_t.grad.numpy()).max():.2e}')
forward max diff: 4.44e-16
grad_img max diff:    4.44e-16
grad_kernel max diff: 2.22e-15
grad_bias max diff:   1.11e-15

Filters you already know: Lesson 12's edge detectors

The output of a conv layer before training is nonsense; the point is that gradient descent will find useful filters. To see what a useful filter's output already looks like, apply a filter you designed by hand back in Lesson 12: Sobel's edge kernel. Every trained CNN's first-layer filters typically converge to something visually similar to this — oriented edge and blob detectors — independent of what the network was trained to do.

In [4]:
test_img = np.zeros((1, 1, 40, 40))
test_img[0, 0, 10:30, 10:30] = 1.0

sobel_x = np.array([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=np.float64).reshape(1, 1, 3, 3)
feature_map = conv2d_forward(test_img, sobel_x, np.zeros(1))

fig, axes = plt.subplots(1, 2, figsize=(6, 3.5))
axes[0].imshow(test_img[0, 0], cmap='gray')
axes[0].set_title('Input')
axes[1].imshow(feature_map[0, 0], cmap='gray')
axes[1].set_title('Feature map\n(Sobel filter, as a conv layer)')
for ax in axes:
    ax.axis('off')
plt.tight_layout()
plt.show()
No description has been provided for this image

Pooling: downsampling with a purpose

Max pooling slides a window over the feature map and keeps only the maximum value in each window, shrinking the spatial size (like the pyramids of Lesson 11) while adding a small amount of local translation invariance — a feature detected at slightly different positions within one pooling window still produces the same output.

In [5]:
def maxpool2d(x, size=2):
    N, C, H, W = x.shape
    OH, OW = H // size, W // size
    out = np.zeros((N, C, OH, OW))
    for i in range(OH):
        for j in range(OW):
            patch = x[:, :, i * size:i * size + size, j * size:j * size + size]
            out[:, :, i, j] = patch.max(axis=(2, 3))
    return out

pool_test = rng.normal(size=(1, 1, 8, 8))
mine = maxpool2d(pool_test)
torch_result = F.max_pool2d(torch.tensor(pool_test), 2)
print(f'max pooling max diff vs torch: {np.abs(mine - torch_result.numpy()).max():.2e}')
max pooling max diff vs torch: 0.00e+00

Why translation equivariance actually matters: a generalization test

Build a small classification task: a $16\times16$ image contains either a plus sign or a circle, at a random position. Train on shapes placed only near the center; test on shapes placed only near the corners — positions the model never saw during training. A network that has genuinely learned "what a plus looks like," rather than "which pixels tend to be on for a plus at these particular training positions," should have no trouble with this.

In [6]:
def make_image(shape_type, cx, cy, size=16):
    img = np.zeros((size, size), dtype=np.float32)
    if shape_type == 'plus':
        img[cy - 1:cy + 2, cx - 3:cx + 4] = 1.0
        img[cy - 3:cy + 4, cx - 1:cx + 2] = 1.0
    else:
        yy, xx = np.mgrid[0:size, 0:size]
        img[((xx - cx) ** 2 + (yy - cy) ** 2) <= 9] = 1.0
    return img

def make_dataset(rng_local, n, position_range):
    imgs, labels = [], []
    for _ in range(n):
        shape_type = rng_local.choice(['plus', 'circle'])
        cx, cy = rng_local.integers(*position_range), rng_local.integers(*position_range)
        imgs.append(make_image(shape_type, cx, cy))
        labels.append(0.0 if shape_type == 'plus' else 1.0)
    return np.array(imgs, dtype=np.float32), np.array(labels, dtype=np.float32)

data_rng = np.random.default_rng(1)
X_train, y_train = make_dataset(data_rng, 300, (5, 11))   # near center
X_test, y_test = make_dataset(data_rng, 150, (3, 5))       # near a corner -- never seen in training

fig, axes = plt.subplots(1, 4, figsize=(9, 2.5))
for ax, im, title in zip(axes, [X_train[0], X_train[1], X_test[0], X_test[1]],
                          ['train example', 'train example', 'test example\n(unseen position)', 'test example\n(unseen position)']):
    ax.imshow(im, cmap='gray')
    ax.set_title(title, fontsize=8)
    ax.axis('off')
plt.tight_layout()
plt.show()
No description has been provided for this image
In [7]:
class MLPClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Flatten(), nn.Linear(16 * 16, 32), nn.ReLU(), nn.Linear(32, 1))

    def forward(self, x):
        return self.net(x).squeeze(-1)

class CNNClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(1, 8, 5, padding=2), nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(8, 16, 5, padding=2), nn.ReLU(),
            nn.AdaptiveMaxPool2d(1),  # global max pool: collapses ALL spatial position info
        )
        self.fc = nn.Linear(16, 1)

    def forward(self, x):
        return self.fc(self.conv(x).flatten(1)).squeeze(-1)

def train_and_eval(model_cls, Xtr, ytr, Xte, yte, seed, epochs=400, lr=0.01):
    torch.manual_seed(seed)  # seed BEFORE constructing the model, so init is actually reproducible
    model = model_cls()
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    for _ in range(epochs):
        opt.zero_grad()
        loss = F.binary_cross_entropy_with_logits(model(Xtr), ytr)
        loss.backward()
        opt.step()
    with torch.no_grad():
        train_acc = ((model(Xtr) > 0).float() == ytr).float().mean().item()
        test_acc = ((model(Xte) > 0).float() == yte).float().mean().item()
    return train_acc, test_acc

Xtr_t = torch.tensor(X_train).unsqueeze(1)
ytr_t = torch.tensor(y_train)
Xte_t = torch.tensor(X_test).unsqueeze(1)
yte_t = torch.tensor(y_test)

mlp_train_acc, mlp_test_acc = train_and_eval(MLPClassifier, Xtr_t, ytr_t, Xte_t, yte_t, seed=0)
cnn_train_acc, cnn_test_acc = train_and_eval(CNNClassifier, Xtr_t, ytr_t, Xte_t, yte_t, seed=0)

print(f'{"model":>6} {"train acc":>10} {"test acc (unseen positions)":>30}')
print(f'{"MLP":>6} {mlp_train_acc:>10.1%} {mlp_test_acc:>30.1%}')
print(f'{"CNN":>6} {cnn_train_acc:>10.1%} {cnn_test_acc:>30.1%}')
 model  train acc    test acc (unseen positions)
   MLP     100.0%                          73.3%
   CNN     100.0%                         100.0%

Both models fit the training data perfectly. On positions neither has ever seen, the flatten-based MLP does little better than a coin flip — it memorized which pixels tend to be on for each class at the training positions, and that knowledge doesn't transfer. The CNN, whose global max pooling forces the final decision to depend only on what filters fired somewhere, not where, generalizes to the new positions with no loss in accuracy at all.

Exercise

  1. Replace nn.AdaptiveMaxPool2d(1) in CNNClassifier with nn.Flatten() directly on the conv features (removing the global pooling, so spatial position is preserved all the way to the final linear layer). Retrain and re-evaluate on the unseen-position test set. Does the CNN's generalization advantage survive?
  2. Increase the gap between train and test position ranges (e.g. train on (5, 11), test on (0, 3), right at the image border where shapes get clipped). Does the CNN's accuracy hold up, or does it degrade — and if it degrades, is that a translation-invariance failure or something else entirely (think about what happens to a shape's appearance, not just its position, right at an edge)?
  3. conv2d_forward above pads nothing (OH = H - KH + 1), so the output shrinks with every layer. Modify it to support zero-padding (pad the input by KH//2 on each side before sliding, matching PyTorch's padding= argument) and confirm the output size matches the input size, validated against F.conv2d(..., padding=1).