PT04 · Convolutional neural networks from scratch

In this lab we implement the basic building blocks of convolutional neural networks (convolutional layers, pooling, batch normalization, residual blocks, dropout, ...) from scratch in PyTorch, and we use them to classify cursive Japanese characters. All of these can be found (already implemented) in the torch.nn module, making the lab highly didactical (but very informative on what a convolution really computes, and on how modules work in PyTorch).

The notebook has two parts, one for each lecture:

• Part 1 (after the lecture on CNNs): convolutions, pooling and depthwise convolutions, a complete CNN, and a look at what it learned.
• Part 2 (after the lecture on deeper models): batch normalization, residual blocks and dropout.

Activity	What you write	Section
1	A convolutional layer with unfold	3	★★
2	Max-pooling and global average pooling	4	★
3	A depthwise convolution	5	★★
4	A complete CNN	6	★
5	A test for translation equivariance	7	★★
6	An occlusion map	8	★★
7	Batch normalization	10	★★
8	Fusing batch normalization into a convolution	11	★★
9	Two bugs with batch normalization	12	★★
10	A residual block	13	★
11	Dropout	14	★

Exercises are rated from 1 to 3 stars based on expected difficulty. Five optional sections close the notebook: a locally-connected layer, patches without copies, receptive fields and causal convolutions (after Part 1), and data augmentation with MixUp (after Part 2).

How to work. Run the cells in order. Each activity is followed by a check cell: if your code is wrong, it stops with a message telling you what to look at. A function you have not written yet raises NotImplementedError the first time it is used, so complete the activities in order. Data loading and the training loop are given.

We assume PT01 and PT02 (tensors, autograd and the training loop). Everything runs on the CPU in a few minutes; in Colab, a GPU runtime makes training faster but is not needed.

Try to solve these exercises without AI coding assistants.

1. Setup and data

We use jaxtyping to annotate shapes, as in PT02, and torchvision for loading and preprocessing the data.

%pip install jaxtyping --quiet

import copy
import math
import time
import matplotlib.pyplot as plt
import torch
import torch.nn.functional as F
import torchvision
from torch.nn import Parameter
from torchvision.transforms import v2
from jaxtyping import Float, Int

torch.manual_seed(0)
torch.set_printoptions(precision=3, sci_mode=False)
print("PyTorch", torch.__version__, "· torchvision", torchvision.__version__)

For this lab we use KMNIST, which is a simple drop-in replacement for MNIST. The input is a (28, 28) black-and-white image of a cursive Japanese character (Kuzushiji), which has to be classified into 10 different classes. There are 60,000 training images and 10,000 test images.

train_raw = torchvision.datasets.KMNIST("./kmnist", train=True, download=True)
test_raw = torchvision.datasets.KMNIST("./kmnist", train=False, download=True)
classes = train_raw.classes
print(len(train_raw), "training images,", len(test_raw), "test images")

# A Dataset is fundamentally a sequence of (input, target) pairs.
image, label = train_raw[0]
print(type(image), label, classes[label])

fig, axes = plt.subplots(3, 10, figsize=(10, 3.4))
for c in range(10):
    examples = (train_raw.targets == c).nonzero().squeeze(1)[:3]
    for row, index in enumerate(examples):
        axes[row, c].imshow(train_raw.data[index], cmap="gray_r")
        axes[row, c].axis("off")
    axes[0, c].set_title(classes[c])
fig.tight_layout()
plt.show()

Transformations (see here) are applied when loading each image. Here we transform the PIL images into PyTorch tensors and standardize them, with a mean and standard deviation computed on the training images only. We keep 6,000 training images aside for validation. The test set is used only once, at the end of Part 1.

split = torch.randperm(60_000, generator=torch.Generator().manual_seed(0))
train_idx, val_idx = split[:54_000], split[54_000:]
pixels = train_raw.data[train_idx].float() / 255
mean, std = pixels.mean().item(), pixels.std().item()
print(f"Pixel mean {mean:.4f}, std {std:.4f} (training images)")

transform = v2.Compose([
    v2.ToImage(),
    v2.ToDtype(torch.float32, scale=True),
    v2.Normalize((mean,), (std,)),
])
full_data = torchvision.datasets.KMNIST("./kmnist", train=True, transform=transform)
train_data = torch.utils.data.Subset(full_data, train_idx.tolist())
val_data = torch.utils.data.Subset(full_data, val_idx.tolist())
test_data = torchvision.datasets.KMNIST("./kmnist", train=False, transform=transform)

# Loaders are used to shuffle and batch the elements of the dataset.
train_loader = torch.utils.data.DataLoader(train_data, batch_size=64, shuffle=True)
val_loader = torch.utils.data.DataLoader(val_data, batch_size=500)
test_loader = torch.utils.data.DataLoader(test_data, batch_size=500)

xb, yb = next(iter(train_loader))
# Shape is now (batch_size, channels, height, width), and the targets are (batch_size,).
print(xb.shape, yb.shape)
assert xb.shape == (64, 1, 28, 28) and yb.shape == (64,)
background = (0 - mean) / std  # The value of an empty pixel (0) after standardization.

The training loop is given. It is the mini-batch loop of PT02, with the AdamW optimizer and a validation pass at the end of every epoch. If a GPU is available, fit uses it and moves the model back to the CPU at the end.

device = "cuda" if torch.cuda.is_available() else "cpu"
print("Training on", device)

@torch.no_grad()
def evaluate(model, loader):
    model.eval()
    model_device = next(model.parameters()).device
    total_loss, correct, count = 0.0, 0, 0
    for x, y in loader:
        x, y = x.to(model_device), y.to(model_device)
        logits = model(x)
        total_loss += F.cross_entropy(logits, y, reduction="sum").item()
        correct += (logits.argmax(dim=1) == y).sum().item()
        count += len(y)
    return total_loss / count, correct / count

def fit(model, loader=train_loader, epochs=3, lr=1e-3, verbose=True):
    model.to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
    history = {"train_loss": [], "val_loss": [], "val_accuracy": []}
    for epoch in range(epochs):
        model.train()
        start = time.time()
        for x, y in loader:
            x, y = x.to(device), y.to(device)
            optimizer.zero_grad()
            loss = F.cross_entropy(model(x), y)
            loss.backward()
            optimizer.step()
            history["train_loss"].append(loss.item())
        val_loss, val_accuracy = evaluate(model, val_loader)
        history["val_loss"].append(val_loss)
        history["val_accuracy"].append(val_accuracy)
        if verbose:
            print(f"Epoch {epoch + 1}: validation loss {val_loss:.3f}, "
                  f"accuracy {val_accuracy:.3f} ({time.time() - start:.0f} s)")
    model.cpu()
    return history

def plot_histories(histories):
    fig, axes = plt.subplots(1, 2, figsize=(10, 3.5))
    for name, history in histories.items():
        losses = torch.tensor(history["train_loss"])
        smooth = losses.unfold(0, 50, 50).mean(dim=1)  # Average of every 50 updates.
        axes[0].plot(torch.arange(len(smooth)) * 50, smooth, label=name)
        epochs = range(1, len(history["val_accuracy"]) + 1)
        axes[1].plot(epochs, history["val_accuracy"], "o-", label=name)
    axes[0].set(xlabel="Update", ylabel="Training loss", yscale="log")
    axes[1].set(xlabel="Epoch", ylabel="Validation accuracy")
    for ax in axes:
        ax.legend()
    fig.tight_layout()
    plt.show()

2. Layers are modules

PyTorch has many pre-implemented layers, such as a 2D convolutional layer.

Predict. The layer below has 1 input channel, 16 output channels, a 5×5 kernel, padding 2 and stride 2. What is the shape of conv2d(xb)? How many parameters does it have? Run the next cell to check.

conv2d = torch.nn.Conv2d(1, 16, 5, padding=2, stride=2)
print(conv2d(xb).shape)
for name, parameter in conv2d.named_parameters():
    print(name, tuple(parameter.shape))
print("Parameters:", sum(p.numel() for p in conv2d.parameters()))

Most of these modules are subclasses of the nn.Module class, which specifies their initialization, their forward pass, their trainable (or non-trainable) parameters, etc. As a warm-up, here is a linear layer $f(\mathbf{X}) = \mathbf{X}\mathbf{W} + \mathbf{b}$ written from scratch. We will use it in the rest of the lab.

Two things to notice. Wrapping a tensor in a Parameter is what makes it appear in model.parameters(), so the optimizer can find it. And we choose the initialization ourselves: the weights are Gaussian with variance $2/d$ for $d$ inputs (He initialization), which keeps the scale of the activations roughly constant through ReLU layers.

class MyLinear(torch.nn.Module):
    def __init__(self, in_features: int, out_features: int):
        super().__init__()
        self.weight = Parameter(torch.randn(in_features, out_features) * math.sqrt(2 / in_features))
        self.bias = Parameter(torch.zeros(out_features))

    def forward(self, xb: Float[torch.Tensor, "batch in_features"]) -> Float[torch.Tensor, "batch out_features"]:
        return xb @ self.weight + self.bias

linear = MyLinear(16, 32)
print(linear)
print([(name, tuple(p.shape)) for name, p in linear.named_parameters()])
assert linear(torch.randn(8, 16)).shape == (8, 32)

Part 1 · Convolutional networks

3. Activity 1: a convolutional layer ★★

Now let us re-implement the Conv2d module. In general, there are many ways to implement a convolution (e.g., with FFTs). In the lecture we saw a simple one: collect the flattened, padded $s \times s$ neighborhoods of every pixel as the columns of a matrix $U(X)$. The convolution is then a single matrix product:

$$\underbrace{\widetilde{Z}}_{(C',\,HW)} = \underbrace{W^\top}_{(C',\,s^2C)}\, \underbrace{U(X)}_{(s^2C,\,HW)} + \mathbf{b}.$$

In PyTorch, unfold builds $U(X)$ for a whole batch: for an input of shape $(B, C, H, W)$ it returns a tensor of shape $(B, C s^2, L)$, where $L$ is the number of sliding windows. Convolution is then matrix multiplication on these patches! Let us look at a tiny example first.

tiny = torch.arange(1., 10.).view(1, 1, 3, 3)
print(tiny)
patches = F.unfold(tiny, kernel_size=2)
print(patches.shape)
print(patches)

There are four 2×2 windows, one per column, read row by row. Each column contains the four pixels of its window, again row by row. With more channels, the rows of each column are ordered by channel first, then by position in the window: this is the same order you get if you flatten a filter of shape $(C, s, s)$.

Predict. What is the shape of F.unfold(xb, kernel_size=3, padding=1)?

print(F.unfold(xb, kernel_size=3, padding=1).shape)

TODO (★★): using unfold, complete the layer below without any other tool from torch.nn apart from Parameter.

• weight must have shape (out_channels, in_channels, kernel_size, kernel_size) and bias shape (out_channels,), as in torch.nn.Conv2d: the checks copy your weights into PyTorch's function to compare the results. Initialize the weights with He initialization (here $d = C s^2$ inputs per output) and the bias with zeros.
• Support zero-padding and a stride. The output has height $\lfloor (H + 2p - s)/\text{stride} \rfloor + 1$, and similarly for the width.

Hint 1: self.weight.view(out_channels, -1) flattens each filter in the same order used by unfold.

Hint 2: a matrix of shape $(C', s^2C)$ times a tensor of shape $(B, s^2C, L)$ broadcasts over the batch. After that, you only need to reshape $L$ back into the output height and width.

class MyConv2d(torch.nn.Module):
    def __init__(self, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, padding: int = 0):
        super().__init__()
        # TODO: complete the code.
        raise NotImplementedError("Complete Activity 1: MyConv2d.__init__")

    def forward(self, xb: Float[torch.Tensor, "batch in_channels height width"]) -> Float[torch.Tensor, "batch out_channels new_height new_width"]:
        # TODO: complete the code.
        raise NotImplementedError("Complete Activity 1: MyConv2d.forward")

def check_conv():
    gen = torch.Generator().manual_seed(1)
    configs = [  # (batch, in, out, height, width, kernel, stride, padding)
        (2, 1, 4, 6, 6, 3, 1, 0), (2, 3, 5, 7, 10, 3, 1, 1), (3, 2, 4, 9, 8, 5, 2, 2),
        (1, 4, 3, 6, 5, 1, 1, 0), (2, 2, 3, 11, 9, 3, 2, 0),
    ]
    for B, C, C_out, H, W, s, stride, p in configs:
        layer = MyConv2d(C, C_out, s, stride=stride, padding=p)
        params = dict(layer.named_parameters())
        assert set(params) == {"weight", "bias"}, f"Expected parameters weight and bias, found {sorted(params)}"
        assert layer.weight.shape == (C_out, C, s, s), f"weight has shape {tuple(layer.weight.shape)}"
        assert layer.bias.shape == (C_out,), f"bias has shape {tuple(layer.bias.shape)}"
        with torch.no_grad():
            layer.bias.normal_(generator=gen)  # A zero bias would hide a broadcasting mistake.
        x = torch.randn(B, C, H, W, generator=gen, requires_grad=True)
        out = layer(x)
        expected = F.conv2d(x, layer.weight, layer.bias, stride=stride, padding=p)
        assert out.shape == expected.shape, f"{(B, C, H, W)} with kernel {s}, stride {stride}, padding {p}: output {tuple(out.shape)}, expected {tuple(expected.shape)}"
        torch.testing.assert_close(out, expected, msg=f"Wrong values for input {(B, C, H, W)}, kernel {s}, stride {stride}, padding {p}")
        seed = torch.randn(out.shape, generator=gen)
        grads = torch.autograd.grad((out * seed).sum(), (x, layer.weight, layer.bias))
        expected_grads = torch.autograd.grad((expected * seed).sum(), (x, layer.weight, layer.bias))
        for g, e, name in zip(grads, expected_grads, ["input", "weight", "bias"]):
            torch.testing.assert_close(g, e, msg=f"Wrong gradient for the {name}")
    wide = MyConv2d(64, 32, 3)
    std = wide.weight.std().item()
    assert abs(std - math.sqrt(2 / (64 * 9))) < 0.01, f"Weight std is {std:.4f}, expected about {math.sqrt(2 / 576):.4f} (He initialization)"
    assert torch.all(wide.bias == 0), "Initialize the bias with zeros"

check_conv()
conv = MyConv2d(1, 16, 3, padding=1)
ypred = conv(xb)
assert ypred.shape == (64, 16, 28, 28)
conv = MyConv2d(16, 32, 5, stride=2, padding=2)
assert conv(ypred).shape == (64, 32, 14, 14)
print("Activity 1 checks passed.")

4. Activity 2: pooling ★

We now move to max-pooling. With your knowledge of unfold, this should be much easier.

TODO (★): complete the custom implementation of max-pooling below, with non-overlapping windows of size kernel_size (as torch.nn.MaxPool2d does by default). If the size of the image is not a multiple of the window, ignore the last rows and columns, as PyTorch does. Then complete global_avg_pool, which averages each channel over all positions, from $(B, C, H, W)$ to $(B, C)$.

Hint: exploit the stride parameter of unfold to avoid overlap between the patches, and then apply the pooling operation on each patch. Be careful with the channels: unlike in a convolution, they are not mixed.

class MyMaxPool2d(torch.nn.Module):
    def __init__(self, kernel_size: int):
        super().__init__()
        self.kernel_size = kernel_size

    def forward(self, xb: Float[torch.Tensor, "batch channels height width"]) -> Float[torch.Tensor, "batch channels reduced_height reduced_width"]:
        # TODO: complete the code. Note: reduced_height = height // kernel_size, and similarly for the width.
        raise NotImplementedError("Complete Activity 2: MyMaxPool2d.forward")

def global_avg_pool(xb: Float[torch.Tensor, "batch channels height width"]) -> Float[torch.Tensor, "batch channels"]:
    # TODO: complete the code.
    raise NotImplementedError("Complete Activity 2: global_avg_pool")

def check_pool():
    gen = torch.Generator().manual_seed(2)
    for shape, k in [((2, 3, 8, 8), 2), ((2, 3, 7, 9), 2), ((1, 2, 9, 6), 3), ((3, 4, 28, 28), 2)]:
        x = (torch.randn(shape, generator=gen) - 2).requires_grad_()  # Mostly negative values.
        out = MyMaxPool2d(k)(x)
        expected = F.max_pool2d(x, k)
        assert out.shape == expected.shape, f"Input {shape}, window {k}: output {tuple(out.shape)}, expected {tuple(expected.shape)}"
        torch.testing.assert_close(out, expected, msg=f"Wrong values for input {shape} and window {k}")
        seed = torch.randn(out.shape, generator=gen)
        torch.testing.assert_close(torch.autograd.grad((out * seed).sum(), x)[0],
                                   torch.autograd.grad((expected * seed).sum(), x)[0], msg="Wrong gradient")
    x = torch.randn(2, 3, 5, 7, generator=gen)
    assert global_avg_pool(x).shape == (2, 3), f"global_avg_pool returned shape {tuple(global_avg_pool(x).shape)}"
    torch.testing.assert_close(global_avg_pool(x), F.adaptive_avg_pool2d(x, 1).flatten(1))

check_pool()
pool = MyMaxPool2d(2)
ypred = pool(xb)
assert ypred.shape == (64, 1, 14, 14)
assert pool(pool(ypred)).shape == (64, 1, 3, 3)
print("Activity 2 checks passed.")

Predict. Max-pooling has no parameters, but we still need to backpropagate through it. If we call backward on the sum of the pooled values, which entries of the input receive a nonzero gradient, and what is their value?

x = torch.randn(1, 1, 4, 4, generator=torch.Generator().manual_seed(3), requires_grad=True)
MyMaxPool2d(2)(x).sum().backward()
print(x.detach().squeeze())
print(x.grad.squeeze())

5. Activity 3: depthwise convolutions ★★

Remember from the lecture that a convolution mixes information in two ways: across positions (spatial mixing) and across channels (channel mixing). A $1 \times 1$ convolution only mixes channels: it is the positionwise MLP we saw in the MLP lecture. A depthwise convolution does the opposite: each channel has its own spatial filter, and channels are never mixed. For a filter $K \sim (s, s, C)$ with $s = 2k + 1$:

$$Y_{i,j,a} = \sum_{u=-k}^{k}\sum_{v=-k}^{k} K_{u+k,\,v+k,\,a}\, X_{i+u,\,j+v,\,a} + b_a.$$

A depthwise convolution followed by a $1 \times 1$ one is a depthwise-separable convolution. In PyTorch, a depthwise convolution is torch.nn.Conv2d(C, C, s, groups=C).

TODO (★★): complete MyDepthwiseConv2d (stride 1). As in PyTorch, weight has shape (channels, 1, kernel_size, kernel_size) and bias shape (channels,). Use He initialization with $d = s^2$ inputs per output.

Hint: the rows of unfold are ordered by channel first, so you can view its output as $(B, C, s^2, L)$. Each channel now has its own patches: you do not need a matrix product.

class MyDepthwiseConv2d(torch.nn.Module):
    def __init__(self, channels: int, kernel_size: int, padding: int = 0):
        super().__init__()
        # TODO: complete the code.
        raise NotImplementedError("Complete Activity 3: MyDepthwiseConv2d.__init__")

    def forward(self, xb: Float[torch.Tensor, "batch channels height width"]) -> Float[torch.Tensor, "batch channels new_height new_width"]:
        # TODO: complete the code.
        raise NotImplementedError("Complete Activity 3: MyDepthwiseConv2d.forward")

def check_depthwise():
    gen = torch.Generator().manual_seed(4)
    for B, C, H, W, s, p in [(2, 3, 6, 6, 3, 0), (2, 4, 7, 9, 3, 1), (1, 2, 8, 5, 5, 2)]:
        layer = MyDepthwiseConv2d(C, s, padding=p)
        assert layer.weight.shape == (C, 1, s, s) and layer.bias.shape == (C,), "Check the shapes of weight and bias"
        with torch.no_grad():
            layer.bias.normal_(generator=gen)
        x = torch.randn(B, C, H, W, generator=gen, requires_grad=True)
        out = layer(x)
        expected = F.conv2d(x, layer.weight, layer.bias, padding=p, groups=C)
        assert out.shape == expected.shape, f"Output {tuple(out.shape)}, expected {tuple(expected.shape)}"
        torch.testing.assert_close(out, expected, msg=f"Wrong values for input {(B, C, H, W)} and kernel {s}")
        seed = torch.randn(out.shape, generator=gen)
        torch.testing.assert_close(torch.autograd.grad((out * seed).sum(), x)[0],
                                   torch.autograd.grad((expected * seed).sum(), x)[0], msg="Wrong gradient")
    std = MyDepthwiseConv2d(512, 3).weight.std().item()
    assert abs(std - math.sqrt(2 / 9)) < 0.03, f"Weight std is {std:.3f}, expected about {math.sqrt(2 / 9):.3f}"

check_depthwise()
print("Activity 3 checks passed.")

A depthwise-separable layer is now just a composition of the two layers you wrote.

Predict. How many parameters does a standard $3 \times 3$ convolution from 64 to 128 channels have? And a separable one with the same input and output channels? Compute them by hand, then run the cell.

class SeparableConv2d(torch.nn.Module):
    def __init__(self, in_channels: int, out_channels: int, kernel_size: int, padding: int = 0):
        super().__init__()
        self.depthwise = MyDepthwiseConv2d(in_channels, kernel_size, padding=padding)
        self.pointwise = MyConv2d(in_channels, out_channels, 1)

    def forward(self, xb):
        return self.pointwise(self.depthwise(xb))

def count_parameters(model):
    return sum(p.numel() for p in model.parameters())

print("Standard: ", count_parameters(MyConv2d(64, 128, 3, padding=1)))
print("Separable:", count_parameters(SeparableConv2d(64, 128, 3, padding=1)))
assert SeparableConv2d(64, 128, 3, padding=1)(torch.randn(2, 64, 8, 8)).shape == (2, 128, 8, 8)

6. Activity 4: building the network ★

Now for the fun part: assemble the previous blocks to build the final convolutional network. As an example, you could do:

x → [MyConv2d → ReLU → MyMaxPool2d] × 2 → MyConv2d → ReLU → global average pooling → MyLinear

TODO (★): complete the specification of the model. Use only the layers you wrote in this notebook (plus F.relu), keep it below 100,000 parameters, and return one logit per class.

The checks reject layers from torch.nn (apart from containers such as torch.nn.ModuleList or torch.nn.Sequential). Feel free to experiment with the number of layers and channels, or with separable convolutions.

class MyCNN(torch.nn.Module):
    def __init__(self, n_classes: int = 10):
        super().__init__()
        # TODO: complete the code.
        raise NotImplementedError("Complete Activity 4: MyCNN.__init__")

    def forward(self, xb: Float[torch.Tensor, "batch channels height width"]) -> Float[torch.Tensor, "batch n_classes"]:
        # TODO: complete the code.
        raise NotImplementedError("Complete Activity 4: MyCNN.forward")

OUR_LAYERS = (MyConv2d, MyMaxPool2d, MyDepthwiseConv2d, SeparableConv2d, MyLinear)
CONTAINERS = (torch.nn.ModuleList, torch.nn.Sequential, torch.nn.ModuleDict)

def check_cnn():
    torch.manual_seed(0)
    model = MyCNN(10)
    for name, module in model.named_modules():
        if name and not isinstance(module, OUR_LAYERS + CONTAINERS) and type(module).__module__.startswith("torch."):
            raise AssertionError(f"{name} is a {type(module).__name__} from torch.nn: use your own layers")
    n = count_parameters(model)
    assert 0 < n < 100_000, f"The model has {n:,} parameters; keep it below 100,000"
    out = model(xb)
    assert out.shape == (64, 10), f"Output shape {tuple(out.shape)}, expected (64, 10)"
    F.cross_entropy(out, yb).backward()
    missing = [name for name, p in model.named_parameters() if p.grad is None or not p.grad.abs().sum() > 0]
    assert not missing, f"These parameters receive no gradient: {missing}"
    return n

print(f"Activity 4 checks passed: {check_cnn():,} parameters.")

We now train the model, and we compare it with an MLP that receives the flattened image (a single hidden layer with 128 units). With the default settings, each epoch of the CNN takes about 20 seconds on a laptop CPU, and longer in a Colab CPU runtime.

EPOCHS = 4
torch.manual_seed(0)
cnn = MyCNN(10)
cnn_history = fit(cnn, epochs=EPOCHS, lr=2e-3)

class MLP(torch.nn.Module):
    def __init__(self, hidden: int = 128, n_classes: int = 10):
        super().__init__()
        self.hidden, self.out = MyLinear(28 * 28, hidden), MyLinear(hidden, n_classes)

    def forward(self, xb):
        return self.out(F.relu(self.hidden(xb.flatten(start_dim=1))))

torch.manual_seed(0)
mlp = MLP()
mlp_history = fit(mlp, epochs=EPOCHS, lr=2e-3)
plot_histories({"CNN": cnn_history, "MLP": mlp_history})
for name, model in [("CNN", cnn), ("MLP", mlp)]:
    print(f"{name}: {count_parameters(model):>7,} parameters, validation accuracy {evaluate(model, val_loader)[1]:.3f}")

Discuss. Which model is better on the validation set? Keep the answer in mind: we will look at the test set at the end of Part 1.

7. Activity 5: translation equivariance ★★

A convolutional layer is translation equivariant: if we shift the input, the output shifts by the same amount. With zero-padding this cannot hold exactly near the borders, since something has to enter the image from outside. The function shift below moves a batch of images by dy rows and dx columns, filling the empty part with fill.

def shift(x: Float[torch.Tensor, "batch channels height width"], dy: int, dx: int, fill: float = 0.0) -> Float[torch.Tensor, "batch channels height width"]:
    H, W = x.shape[-2:]
    out = torch.full_like(x, fill)
    out[..., max(dy, 0):H + min(dy, 0), max(dx, 0):W + min(dx, 0)] = \
        x[..., max(-dy, 0):H - max(dy, 0), max(-dx, 0):W - max(dx, 0)]
    return out

fig, axes = plt.subplots(1, 3, figsize=(6, 2))
for ax, (dy, dx) in zip(axes, [(0, 0), (3, 0), (-2, 4)]):
    ax.imshow(shift(xb[:1], dy, dx, fill=background)[0, 0], cmap="gray_r")
    ax.set_title(f"dy={dy}, dx={dx}")
    ax.axis("off")
plt.show()

TODO (★★): complete equivariance_error, which shifts the input by (dy, dx), applies layer, and compares the result with the shifted output of the original input. Both outputs are cropped by border pixels on each side before computing the largest absolute difference. Use the default fill=0 in both shifts: it is the same value used by zero-padding. Assume that layer keeps the height and width of its input.

Then use it to answer the questions below the check.

def equivariance_error(layer, x: Float[torch.Tensor, "batch channels height width"], dy: int, dx: int, border: int) -> float:
    # TODO: complete the code. Return a Python float.
    raise NotImplementedError("Complete Activity 5: equivariance_error")

def check_equivariance_error():
    identity = lambda x: x
    x = torch.randn(2, 1, 10, 10, generator=torch.Generator().manual_seed(5))
    assert isinstance(equivariance_error(identity, x, 2, 1, 0), float), "Return a Python float, e.g. with .item()"
    assert equivariance_error(identity, x, 2, 1, 0) == 0, "The identity is equivariant everywhere"
    flip = lambda x: x.flip(-1)
    assert equivariance_error(flip, x, 0, 2, 3) > 0.1, "Flipping the image is not equivariant: the error should be large"
    added = lambda x: x + 1  # The constant enters with the shift: wrong everywhere near the empty region.
    assert equivariance_error(added, x, 3, 0, 0) == 1 and equivariance_error(added, x, 3, 0, 3) == 0, \
        "Crop `border` pixels on every side of both outputs before taking the maximum"

check_equivariance_error()
print("Activity 5 checks passed.")

Use equivariance_error on a batch of KMNIST images (for example xb[:8]) to answer:

1. For cnn.conv1 (a trained 3×3 convolution with padding 1) and a shift of 2 pixels in both directions, what is the smallest border for which the error is zero (up to float precision)? Why that value?
2. Repeat with lambda x: cnn.conv2(F.relu(cnn.conv1(x))), two stacked layers. How does the border change?
3. Max-pooling halves the resolution, so the function cannot be used directly. Instead, compare pool(shift(x, d, d)) with shift(pool(x), d // 2, d // 2) for d = 2 and d = 1, with pool = MyMaxPool2d(2). What happens?

images = xb[:8]
# Your experiments here.

What does this mean for classification? The cell below shifts the validation images by an increasing number of pixels (in a random direction, filling with background) and measures the accuracy of the two models.

Predict. Which model degrades faster?

@torch.no_grad()
def shifted_accuracy(model, x, y, d, seed=0):
    model.eval()
    gen = torch.Generator().manual_seed(seed)
    signs = torch.randint(0, 2, (len(x), 2), generator=gen) * 2 - 1
    moved = torch.stack([shift(img[None], d * int(sy), d * int(sx), fill=background)[0]
                         for img, (sy, sx) in zip(x, signs)])
    return (model(moved).argmax(dim=1) == y).float().mean().item()

val_x, val_y = next(iter(torch.utils.data.DataLoader(val_data, batch_size=2000)))
distances = range(0, 6)
results = {name: [shifted_accuracy(model, val_x, val_y, d) for d in distances]
           for name, model in [("CNN", cnn), ("MLP", mlp)]}
plt.figure(figsize=(5, 3.5))
for name, accuracies in results.items():
    plt.plot(distances, accuracies, "o-", label=name)
plt.xlabel("Shift (pixels in each direction)")
plt.ylabel("Validation accuracy")
plt.legend()
plt.show()

8. Activity 6: what did the network learn? ★★

Let us look inside the trained CNN. The first row below shows the filters of the first layer, the second row the corresponding feature maps (after the ReLU) for one validation image.

with torch.no_grad():
    filters = cnn.conv1.weight[:, 0].detach()
    maps = F.relu(cnn.conv1(val_x[:1]))[0]
fig, axes = plt.subplots(2, 16, figsize=(12, 1.8))
for i in range(16):
    axes[0, i].imshow(filters[i], cmap="RdBu_r", vmin=-filters.abs().max(), vmax=filters.abs().max())
    axes[1, i].imshow(maps[i], cmap="gray_r")
for ax in axes.flat:
    ax.axis("off")
plt.show()

The 3×3 filters are hard to read on their own, but the feature maps show what they do: each one responds to the edges or strokes of the character in a different direction. Deeper layers are harder to visualize in this way. A simple alternative is an occlusion map (Zeiler and Fergus, 2014): cover a small square of the image with background, and see how much the probability of the correct class drops. If it drops a lot, the network was using that region.

TODO (★★): complete occlusion_map. For every position of a patch × patch square, with top-left corners every stride pixels (range(0, 28 - patch + 1, stride) in both directions), replace the square with background and compute the probability that model assigns to label. Return a tensor of probabilities with one row per vertical position and one column per horizontal position.

Hint: build all the occluded images in a single batch and call the model once. Put the model in inference mode and do not track gradients.

def occlusion_map(model, image: Float[torch.Tensor, "1 height width"], label: int, patch: int = 7, stride: int = 3) -> Float[torch.Tensor, "rows cols"]:
    # TODO: complete the code.
    raise NotImplementedError("Complete Activity 6: occlusion_map")

def check_occlusion():
    image, label = val_x[0], int(val_y[0])
    result = occlusion_map(cnn, image, label)
    assert result.shape == (8, 8), f"Expected shape (8, 8) for patch 7 and stride 3, found {tuple(result.shape)}"
    assert not result.requires_grad, "Do not track gradients"
    covered = image.clone()
    covered[:, 6:13, 3:10] = background
    with torch.no_grad():
        expected = cnn(covered[None]).softmax(dim=1)[0, label]
    torch.testing.assert_close(result[2, 1], expected, msg="Row 2, column 1 should cover rows 6 to 12 and columns 3 to 9")
    assert occlusion_map(cnn, image, label, patch=4, stride=8).shape == (4, 4)
    assert not cnn.training, "Put the model in inference mode"

check_occlusion()
print("Activity 6 checks passed.")

fig, axes = plt.subplots(2, 5, figsize=(10, 4.2))
for n in range(5):
    image, label = val_x[n], int(val_y[n])
    heat = occlusion_map(cnn, image, label, patch=7, stride=1)
    axes[0, n].imshow(image[0], cmap="gray_r")
    axes[0, n].set_title(classes[label])
    axes[1, n].imshow(heat, cmap="magma", vmin=0, vmax=1)
    axes[1, n].set_title(f"min p = {heat.min():.2f}")
for ax in axes.flat:
    ax.axis("off")
fig.tight_layout()
plt.show()

Discuss. Dark regions are where an occlusion hurts the prediction. Are they on the strokes of the character, or somewhere else? Is the network relying on a single detail, or on the whole shape?

9. The test set

We chose all the settings above by looking at the validation set, so the test set gives an unbiased estimate for the final models.

for name, model in [("CNN", cnn), ("MLP", mlp)]:
    print(f"{name}: validation accuracy {evaluate(model, val_loader)[1]:.3f}, test accuracy {evaluate(model, test_loader)[1]:.3f}")

Discuss. Both models lose accuracy on the test set, but not by the same amount. What does this suggest about the difference between the training and test images, and about which model generalizes better?

Part 2 · Training deeper networks

This part follows the lecture on deeper models. We implement batch normalization and dropout from scratch, and we build residual blocks. To keep the experiments fast, after checking that our layers compute the same thing as PyTorch's, we use torch.nn layers when training larger models.

10. Activity 7: batch normalization ★★

After a convolution, a batch normalization layer standardizes each channel using every position of every image in the mini-batch:

$$\widetilde{\mu}_c = \frac{1}{BHW}\sum_{b,i,j} Z_{b,c,i,j}, \qquad \widetilde{\sigma}_c^2 = \frac{1}{BHW}\sum_{b,i,j} \left(Z_{b,c,i,j} - \widetilde{\mu}_c\right)^2,$$

$$\text{BN}(Z)_{b,c,i,j} = \gamma_c\, \frac{Z_{b,c,i,j} - \widetilde{\mu}_c}{\sqrt{\widetilde{\sigma}_c^2 + \varepsilon}} + \beta_c,$$

where $\gamma_c, \beta_c$ are trainable and $\varepsilon$ is a small value to avoid division by zero. During inference, the mean and the variance are fixed. In the simplest case, we keep a running average during training:

$$\mu_c \leftarrow (1 - m)\,\mu_c + m\,\widetilde{\mu}_c,$$

and similarly for $\sigma_c^2$, where the running variance uses the unbiased estimate of the mini-batch variance (dividing by $BHW - 1$). The momentum $m$ is 0.1 in PyTorch. Thus, for an input with $C$ channels, the layer has $2C$ trainable parameters and $2C$ untrainable ones.

TODO (★★): implement the layer. As in PyTorch, call the parameters weight ($\gamma$, initialized to ones) and bias ($\beta$, zeros), and the running statistics running_mean (zeros) and running_var (ones).

Hint 1: non-trainable state in PyTorch should be handled with buffers: they are saved with the model and moved with .to(), but they are not returned by .parameters().

Hint 2: the updates of the running statistics should not be differentiated, so they go within a torch.no_grad() scope to avoid "cluttering" the computational graph.

Hint 3: training and inference mode can be checked with the internal self.training flag.

class MyBatchNorm2d(torch.nn.Module):
    def __init__(self, channels: int, eps: float = 1e-5, momentum: float = 0.1):
        super().__init__()
        # TODO: complete the code.
        raise NotImplementedError("Complete Activity 7: MyBatchNorm2d.__init__")

    def forward(self, xb: Float[torch.Tensor, "batch channels height width"]) -> Float[torch.Tensor, "batch channels height width"]:
        # TODO: complete the code.
        raise NotImplementedError("Complete Activity 7: MyBatchNorm2d.forward")

def check_batchnorm():
    gen = torch.Generator().manual_seed(6)
    bn = MyBatchNorm2d(3)
    assert sorted(n for n, _ in bn.named_parameters()) == ["bias", "weight"], "Trainable parameters: weight and bias"
    assert sorted(n for n, _ in bn.named_buffers()) == ["running_mean", "running_var"], "Buffers: running_mean and running_var"
    torch.testing.assert_close(bn.weight, torch.ones(3)); torch.testing.assert_close(bn.bias, torch.zeros(3))
    torch.testing.assert_close(bn.running_mean, torch.zeros(3)); torch.testing.assert_close(bn.running_var, torch.ones(3))
    with torch.no_grad():
        bn.weight.uniform_(0.5, 2., generator=gen); bn.bias.normal_(generator=gen)
    reference_mean, reference_var = torch.zeros(3), torch.ones(3)
    scale = torch.tensor([1., 3., 0.5])[None, :, None, None]
    offset = torch.tensor([2., -1., 5.])[None, :, None, None]
    bn.train()
    for step in range(3):
        x = (torch.randn(4, 3, 5, 6, generator=gen) * scale + offset).requires_grad_()
        out = bn(x)
        expected = F.batch_norm(x, reference_mean, reference_var, bn.weight, bn.bias, training=True, momentum=0.1, eps=1e-5)
        torch.testing.assert_close(out, expected, msg="Wrong output in training mode: use the statistics of the mini-batch")
        seed = torch.randn(out.shape, generator=gen)
        for wrt, name in [(x, "input"), (bn.weight, "weight"), (bn.bias, "bias")]:
            torch.testing.assert_close(torch.autograd.grad((out * seed).sum(), wrt, retain_graph=True)[0],
                                       torch.autograd.grad((expected * seed).sum(), wrt, retain_graph=True)[0],
                                       msg=f"Wrong gradient with respect to the {name}")
        torch.testing.assert_close(bn.running_mean, reference_mean, msg=f"Wrong running_mean after {step + 1} training steps")
        torch.testing.assert_close(bn.running_var, reference_var, msg=f"Wrong running_var after {step + 1} training steps (unbiased estimate, momentum 0.1)")
    assert not bn.running_mean.requires_grad and bn.running_mean.grad_fn is None, "Update the running statistics without tracking gradients"
    bn.eval()
    x = torch.randn(2, 3, 4, 4, generator=gen) * scale + offset
    before = bn.running_mean.clone()
    torch.testing.assert_close(bn(x), F.batch_norm(x, reference_mean, reference_var, bn.weight, bn.bias, training=False, eps=1e-5),
                               msg="Wrong output in inference mode: use the running statistics")
    torch.testing.assert_close(bn.running_mean, before, msg="Inference must not change the running statistics")

check_batchnorm()
bn = MyBatchNorm2d(1)
out = bn(xb)
assert out.shape == (64, 1, 28, 28)
torch.testing.assert_close(out.mean(dim=(0, 2, 3)), torch.zeros(1), atol=1e-5, rtol=0)
torch.testing.assert_close(out.std(dim=(0, 2, 3)), torch.ones(1), atol=1e-2, rtol=0)
print("Activity 7 checks passed.")

11. Activity 8: fusing batch normalization ★★

At inference, the statistics are fixed, so batch normalization is just an affine function of each channel:

$$\text{BN}(z) = a \odot z + c, \qquad a_c = \frac{\gamma_c}{\sqrt{\sigma_c^2 + \varepsilon}}, \qquad c_c = \beta_c - a_c\,\mu_c.$$

Composed with the convolution before it, it gives another convolution: each filter is multiplied by $a_c$, and the new bias is $a_c b_c + c_c$. Deployed networks do this to save one layer.

TODO (★★): complete fuse_conv_bn, which returns a new MyConv2d that computes bn(conv(x)) in inference mode. Do not modify conv or bn.

Hint: start from copy.deepcopy(conv) and overwrite its parameters within torch.no_grad().

def fuse_conv_bn(conv: MyConv2d, bn: MyBatchNorm2d) -> MyConv2d:
    # TODO: complete the code.
    raise NotImplementedError("Complete Activity 8: fuse_conv_bn")

def check_fusion():
    gen = torch.Generator().manual_seed(7)
    for stride, padding in [(1, 1), (2, 0)]:
        conv, bn = MyConv2d(3, 5, 3, stride=stride, padding=padding), MyBatchNorm2d(5)
        with torch.no_grad():
            conv.bias.normal_(generator=gen)
            bn.weight.uniform_(0.5, 2., generator=gen); bn.bias.normal_(generator=gen)
            bn.running_mean.normal_(generator=gen); bn.running_var.uniform_(0.5, 3., generator=gen)
        saved = [t.clone() for t in list(conv.parameters()) + list(bn.parameters()) + list(bn.buffers())]
        bn.eval()
        fused = fuse_conv_bn(conv, bn)
        assert isinstance(fused, MyConv2d) and fused is not conv, "Return a new MyConv2d"
        x = torch.randn(2, 3, 9, 9, generator=gen)
        with torch.no_grad():
            torch.testing.assert_close(fused(x), bn(conv(x)), msg=f"The fused layer differs from bn(conv(x)) (stride {stride}, padding {padding})")
        current = list(conv.parameters()) + list(bn.parameters()) + list(bn.buffers())
        assert all(torch.equal(a, b) for a, b in zip(saved, current)), "fuse_conv_bn modified conv or bn"

check_fusion()
print("Activity 8 checks passed.")

Discuss. Why can we fuse the two layers only at inference, and not during training?

12. Activity 9: two bugs with batch normalization ★★

Batch normalization is easy to get wrong, because it has two modes and a hidden state. We first train a small CNN with batch normalization (using PyTorch's layers, which compute the same thing as yours). Then we look at two pieces of code that have a bug.

def bn_cnn():
    def block(c_in, c_out):
        return [torch.nn.Conv2d(c_in, c_out, 3, padding=1, bias=False), torch.nn.BatchNorm2d(c_out), torch.nn.ReLU()]
    return torch.nn.Sequential(*block(1, 32), torch.nn.MaxPool2d(2), *block(32, 64), torch.nn.MaxPool2d(2),
                               *block(64, 64), torch.nn.AdaptiveAvgPool2d(1), torch.nn.Flatten(), torch.nn.Linear(64, 10))

torch.manual_seed(0)
model_bn = bn_cnn()
_ = fit(model_bn, epochs=1, lr=2e-3)

Case 1. Is the prediction a property of the image? A colleague fine-tuned the model for one more step on a new batch of data, and then used it to classify a few images. They noticed that the probability given to the same image changes when it is classified together with other images.

case1_model = copy.deepcopy(model_bn)
optimizer = torch.optim.AdamW(case1_model.parameters(), lr=1e-4)
case1_model.train()
x, y = next(iter(train_loader))
optimizer.zero_grad()
F.cross_entropy(case1_model(x), y).backward()
optimizer.step()

@torch.no_grad()
def predict(model, images):
    return model(images).softmax(dim=1)

image, label = val_x[:1], int(val_y[0])
same_class = val_x[val_y == label][1:8]
other_classes = val_x[val_y != label][:7]
print("Alone:                    ", predict(case1_model, image)[0, label].item())
print("With 7 of the same class: ", predict(case1_model, torch.cat([image, same_class]))[0, label].item())
print("With 7 of other classes:  ", predict(case1_model, torch.cat([image, other_classes]))[0, label].item())

Case 2. Who changed my model? A second colleague wanted to quickly check how the model handles images with inverted colors (dark characters on a light background, the opposite of KMNIST), and wrote a short loop under torch.no_grad(). After that, the accuracy on the ordinary validation set is lower than before.

case2_model = copy.deepcopy(model_bn)
print("Validation accuracy before:", evaluate(case2_model, val_loader)[1])
with torch.no_grad():
    case2_model.train()
    for x, y in torch.utils.data.DataLoader(val_data, batch_size=500):
        inverted = (1 - 2 * mean) / std - x  # The standardized version of 1 - pixel.
        accuracy = (case2_model(inverted).argmax(dim=1) == y).float().mean()
print("Validation accuracy after: ", evaluate(case2_model, val_loader)[1])

TODO (★★): explain what goes wrong in each case. Then write safe_evaluate(model, x), which returns the predicted probabilities for a batch x in a way that avoids both bugs. It must (1) use inference mode, (2) check with an assert that the call did not change any buffer of the model, and (3) leave the model in the mode (training or inference) it was in before the call.

def safe_evaluate(model, x: Float[torch.Tensor, "batch channels height width"]) -> Float[torch.Tensor, "batch n_classes"]:
    # TODO: complete the code.
    raise NotImplementedError("Complete Activity 9: safe_evaluate")

def check_safe_evaluate():
    model = copy.deepcopy(model_bn)
    for mode in [True, False]:
        model.train(mode)
        state = copy.deepcopy(model.state_dict())
        alone = safe_evaluate(model, val_x[:1])
        together = safe_evaluate(model, val_x[:64])[:1]
        torch.testing.assert_close(alone, together, msg="The prediction for an image must not depend on the rest of the batch")
        assert model.training == mode, "Restore the previous mode of the model"
        assert all(torch.equal(state[k], v) for k, v in model.state_dict().items()), "safe_evaluate changed the model"
    assert torch.allclose(alone.sum(dim=1), torch.ones(1)), "Return probabilities"

    class Sneaky(torch.nn.Module):  # A model whose forward always changes its buffer.
        def __init__(self):
            super().__init__()
            self.register_buffer("calls", torch.zeros(()))
        def forward(self, x):
            self.calls += 1
            return x.flatten(1)[:, :10]
    try:
        safe_evaluate(Sneaky(), val_x[:2])
    except AssertionError:
        pass
    else:
        raise AssertionError("safe_evaluate should fail when the call changes a buffer")

check_safe_evaluate()
print("Activity 9 checks passed.")

13. Activity 10: residual blocks ★

In the lecture we saw that very deep plain networks are hard to optimize, and that residual connections fix this. The block we use is the one from the lecture: two 3×3 convolutions with batch normalization, a ReLU between them, the sum with the input, and a final ReLU after the sum:

$$\mathbf{h} = \text{ReLU}\big(\mathbf{x} + f(\mathbf{x})\big), \qquad f(\mathbf{x}) = \text{BN}(\text{Conv}(\text{ReLU}(\text{BN}(\text{Conv}(\mathbf{x}))))).$$

The convolutions have no bias, since BN would remove it anyway.

TODO (★): complete ResidualBlock, using torch.nn.Conv2d and torch.nn.BatchNorm2d for speed. With residual=False, the same block computes $\text{ReLU}(f(\mathbf{x}))$, which we use as the plain baseline. With batchnorm=False, replace the BN layers with torch.nn.Identity() and keep the bias of the convolutions.

class ResidualBlock(torch.nn.Module):
    def __init__(self, channels: int, residual: bool = True, batchnorm: bool = True):
        super().__init__()
        # TODO: complete the code.
        raise NotImplementedError("Complete Activity 10: ResidualBlock.__init__")

    def forward(self, xb: Float[torch.Tensor, "batch channels height width"]) -> Float[torch.Tensor, "batch channels height width"]:
        # TODO: complete the code.
        raise NotImplementedError("Complete Activity 10: ResidualBlock.forward")

def check_residual():
    gen = torch.Generator().manual_seed(8)
    x = torch.randn(4, 8, 6, 6, generator=gen)
    torch.manual_seed(0)
    block = ResidualBlock(8)
    assert block(x).shape == x.shape
    n = count_parameters(block)
    assert n == 2 * 8 * 8 * 9 + 4 * 8, f"{n} parameters: two 3x3 convolutions without bias and two BN layers expected"
    batchnorms = [m for m in block.modules() if isinstance(m, torch.nn.BatchNorm2d)]
    assert len(batchnorms) == 2, "Use two BatchNorm2d layers"
    with torch.no_grad():
        batchnorms[-1].weight.zero_()  # Now f(x) = 0.
    block.eval()
    torch.testing.assert_close(block(x), F.relu(x), msg="With f(x) = 0, a residual block should return ReLU(x)")
    torch.manual_seed(0)
    plain = ResidualBlock(8, residual=False)
    with torch.no_grad():
        [m for m in plain.modules() if isinstance(m, torch.nn.BatchNorm2d)][-1].weight.zero_()
    plain.eval()
    torch.testing.assert_close(plain(x), torch.zeros_like(x), msg="Without the residual, the block returns ReLU(f(x))")
    no_bn = ResidualBlock(8, batchnorm=False)
    assert not any(isinstance(m, torch.nn.BatchNorm2d) for m in no_bn.modules()), "batchnorm=False should not use BN"
    assert count_parameters(no_bn) == 2 * (8 * 8 * 9 + 8), "Without BN, the convolutions keep their bias"
    block.train()
    assert block(x).shape == x.shape

check_residual()
print("Activity 10 checks passed.")

We now repeat the opening experiment of the lecture on a small scale. Each network has a first convolution with stride 2 (to 16 channels, 14×14), a stack of blocks, global average pooling and a linear layer. We train it on 10,000 training images for two epochs, for three depths and three variants: plain, plain with BN, and residual with BN. This takes a few minutes on a CPU.

Predict. Before running: how will the training loss of the three variants change with depth?

DEPTHS, SUBSET = [2, 8, 16], 10_000

def deep_cnn(n_blocks, residual, batchnorm, channels=16):
    return torch.nn.Sequential(
        torch.nn.Conv2d(1, channels, 3, stride=2, padding=1), torch.nn.ReLU(),
        *[ResidualBlock(channels, residual, batchnorm) for _ in range(n_blocks)],
        torch.nn.AdaptiveAvgPool2d(1), torch.nn.Flatten(), torch.nn.Linear(channels, 10))

subset_loader = torch.utils.data.DataLoader(torch.utils.data.Subset(train_data, range(SUBSET)), batch_size=64, shuffle=True)
train_check = next(iter(torch.utils.data.DataLoader(torch.utils.data.Subset(train_data, range(2000)), batch_size=2000)))
variants = {"Plain": (False, False), "Plain + BN": (False, True), "Residual + BN": (True, True)}
depth_results = {name: {"train_loss": [], "val_accuracy": []} for name in variants}
for n_blocks in DEPTHS:
    for name, (residual, batchnorm) in variants.items():
        torch.manual_seed(0)
        model = deep_cnn(n_blocks, residual, batchnorm)
        start = time.time()
        history = fit(model, loader=subset_loader, epochs=2, lr=1e-3, verbose=False)
        with torch.no_grad():
            model.eval()
            train_loss = F.cross_entropy(model(train_check[0]), train_check[1]).item()
        depth_results[name]["train_loss"].append(train_loss)
        depth_results[name]["val_accuracy"].append(history["val_accuracy"][-1])
        print(f"{2 * n_blocks + 2:>2} layers, {name:<13}: training loss {train_loss:.3f}, "
              f"validation accuracy {history['val_accuracy'][-1]:.3f} ({time.time() - start:.0f} s)")

layers = [2 * n + 2 for n in DEPTHS]
fig, axes = plt.subplots(1, 2, figsize=(10, 3.5))
for name, result in depth_results.items():
    axes[0].plot(layers, result["train_loss"], "o-", label=name)
    axes[1].plot(layers, result["val_accuracy"], "o-", label=name)
axes[0].axhline(math.log(10), color="gray", linestyle="--", label="Chance")
axes[0].set(xlabel="Layers", ylabel="Training loss")
axes[1].set(xlabel="Layers", ylabel="Validation accuracy", ylim=(0, 1))
fig.legend(*axes[0].get_legend_handles_labels(), loc="upper center", ncol=4, bbox_to_anchor=(0.5, 1.08))
fig.tight_layout()
plt.show()

14. Activity 11: dropout ★

Recall that dropout during training works as:

$$\text{Dropout}(X) = \frac{X \odot M}{1 - p},$$

where $M$ is a binary mask whose entries are 0 with probability $p$, sampled again at every call. During inference, the layer does nothing and leaves the input untouched. As we saw in the lecture, dropout is used much less in modern CNNs and LLMs than it used to be, but it is still used, for instance when fine-tuning large models.

TODO (★): implement the dropout layer below.

Hint 1: if $u$ is sampled according to a uniform distribution, $u > p$ will be 0 with probability $p$.

Hint 2: as for batch normalization, use the internal training flag.

class MyDropout(torch.nn.Module):
    def __init__(self, p: float = 0.5):
        super().__init__()
        self.p = p

    def forward(self, xb: Float[torch.Tensor, "..."]) -> Float[torch.Tensor, "..."]:
        # TODO: complete the code.
        raise NotImplementedError("Complete Activity 11: MyDropout.forward")

def check_dropout():
    torch.manual_seed(9)
    dropout = MyDropout(0.3)
    dropout.eval()
    x = torch.randn(64, 1, 28, 28)
    assert torch.equal(dropout(x), x), "In inference mode, dropout must return its input unchanged"
    dropout.train()
    ones = torch.ones(200_000, requires_grad=True)
    out = dropout(ones)
    dropped = (out == 0).float().mean().item()
    assert abs(dropped - 0.3) < 0.01, f"{dropped:.1%} of the values were dropped, expected about 30%"
    kept = out[out != 0]
    torch.testing.assert_close(kept, torch.full_like(kept, 1 / 0.7), msg="Kept values must be divided by 1 - p")
    out.sum().backward()
    torch.testing.assert_close(ones.grad, out.detach(), msg="The gradient is the mask divided by 1 - p")
    assert not torch.equal(dropout(torch.ones(1000)), dropout(torch.ones(1000))), "Sample a new mask at every call"
    assert torch.equal(MyDropout(0.)(x), x), "With p = 0, dropout does nothing"
    assert dropout(torch.ones(3, 4, 5)).shape == (3, 4, 5)

check_dropout()
print("Activity 11 checks passed.")

Discuss. Why do we divide by $1 - p$ during training, rather than multiplying by $1 - p$ during inference? Both give the same expected value.

What we built

We wrote from scratch most of the layers of a modern CNN (standard and depthwise convolutions, pooling, batch normalization, dropout) and checked each of them against PyTorch. We also saw why a CNN beats an MLP on new images, where translation equivariance breaks, and why deep plain networks need BN and residual connections to train.

Discuss

1. In Activity 1, unfold creates a tensor $s^2$ times larger than the input. Why is this a problem for large images, and how do libraries avoid it? (See Optional B.)
2. A CNN for 28×28 images with global average pooling can be applied to a 56×56 image without changing any weight. Why? What happens with an MLP?
3. Where in this notebook did a model behave differently in training and inference mode?

Optional A: a locally-connected layer ★★★

After Part 1. Having full control over the convolution allows us to implement several interesting variations. A locally-connected layer is identical to the convolutional one, except that the weights are not tied: each sliding window is multiplied by a different kernel matrix.

TODO (★★★): modify your implementation of MyConv2d to obtain a locally-connected variant (stride 1). It needs to know the input size, to create one kernel per window: the weight has shape (L, out_channels, in_channels * kernel_size**2), with L the number of windows, and the bias has shape (L, out_channels). This increases the number of parameters by a factor of L.

Hint: torch.einsum is the easiest way to multiply each window by its own matrix.

class MyLocallyConnected2d(torch.nn.Module):
    def __init__(self, in_channels: int, out_channels: int, kernel_size: int, height: int, width: int, padding: int = 0):
        super().__init__()
        # TODO: complete the code.
        raise NotImplementedError("Optional exercise: MyLocallyConnected2d.__init__")

    def forward(self, xb: Float[torch.Tensor, "batch in_channels height width"]) -> Float[torch.Tensor, "batch out_channels new_height new_width"]:
        # TODO: complete the code.
        raise NotImplementedError("Optional exercise: MyLocallyConnected2d.forward")

local = MyLocallyConnected2d(2, 3, 3, height=6, width=7, padding=1)
assert local.weight.shape == (42, 3, 18) and local.bias.shape == (42, 3)
x = torch.randn(4, 2, 6, 7)
# With the same kernel in every window, the layer is a convolution.
conv = MyConv2d(2, 3, 3, padding=1)
with torch.no_grad():
    local.weight.copy_(conv.weight.view(1, 3, 18).expand(42, -1, -1))
    conv.bias.normal_()
    local.bias.copy_(conv.bias.expand(42, -1))
torch.testing.assert_close(local(x), conv(x), msg="With tied weights, the layer should match MyConv2d")
# Each window has its own kernel: changing one kernel changes only one output position.
with torch.no_grad():
    local.weight[10].normal_()
changed = (local(x) - conv(x)).abs().sum(dim=(0, 1)).flatten().nonzero().squeeze(1).tolist()
assert changed == [10], f"Changing the kernel of window 10 changed positions {changed}"
print("Parameters, convolution:       ", count_parameters(MyConv2d(1, 16, 3, padding=1)))
print("Parameters, locally connected: ", count_parameters(MyLocallyConnected2d(1, 16, 3, 28, 28, padding=1)))
fresh = MyLocallyConnected2d(1, 16, 3, 28, 28, padding=1)
print("Equivariance error (border 3): ", equivariance_error(fresh, xb[:8], 2, 2, 3))

The last line uses your function from Activity 5: unlike a convolution, a locally-connected layer is not equivariant, even far from the borders. It has the same locality as a convolution but no weight sharing, with 784 times more parameters on KMNIST.

Optional B: patches without copies ★★★

After Part 1. unfold copies every pixel into $s^2$ different columns. In PT01 we saw that a tensor is a view on a block of memory, described by its shape and its strides (how many elements to skip to move along each axis). With the right strides, we can describe all the patches of an image as a view, without copying anything.

TODO (★★★): complete patches_view, which receives a (padded) batch x of shape $(B, C, H, W)$ and returns a view of shape $(B, C, H - s + 1, W - s + 1, s, s)$, where out[b, c, i, j] is the $s \times s$ patch with top-left corner $(i, j)$. Use as_strided and x.stride().

Hint: moving to the next patch to the right and moving to the next pixel to the right inside a patch skip the same number of elements in memory.

def patches_view(x: Float[torch.Tensor, "batch channels height width"], kernel_size: int) -> Float[torch.Tensor, "batch channels rows cols kernel kernel"]:
    # TODO: complete the code.
    raise NotImplementedError("Optional exercise: patches_view")

x = torch.randn(2, 3, 6, 7)
view = patches_view(x, 3)
assert view.shape == (2, 3, 4, 5, 3, 3)
assert view.untyped_storage().data_ptr() == x.untyped_storage().data_ptr(), "The result must be a view of x, not a copy"
torch.testing.assert_close(view[1, 2, 3, 1], x[1, 2, 3:6, 1:4])
# Same values as unfold, in a different layout.
torch.testing.assert_close(view.permute(0, 1, 4, 5, 2, 3).reshape(2, 27, 20), F.unfold(x, 3))
torch.testing.assert_close(patches_view(x[:, 1:], 2), F.unfold(x[:, 1:], 2).view(2, 2, 2, 2, 5, 6).permute(0, 1, 4, 5, 2, 3))
print("Is the view contiguous?", view.is_contiguous())

A convolution is now an einsum over the view. The cell compares the memory used by unfold with the size of the input, and measures the time of three implementations on a batch of 64 images with 32 channels. Timings depend a lot on your hardware: we run each function a few times before measuring, and we report the median.

import torch.utils.benchmark as benchmark

def einsum_conv(x, weight, bias, padding):
    view = patches_view(F.pad(x, (padding,) * 4), weight.shape[-1])
    return torch.einsum("bchwij,ocij->bohw", view, weight) + bias[:, None, None]

x = torch.randn(64, 32, 28, 28)
layer = MyConv2d(32, 64, 3, padding=1)
with torch.no_grad():
    torch.testing.assert_close(einsum_conv(x, layer.weight, layer.bias, 1), layer(x), rtol=1e-4, atol=1e-4)
    unfold_bytes = F.unfold(x, 3, padding=1).numel() * 4
print(f"Input: {x.numel() * 4 / 2**20:.1f} MiB, unfold: {unfold_bytes / 2**20:.1f} MiB")

functions = {
    "MyConv2d (unfold)": lambda: layer(x),
    "einsum on a view": lambda: einsum_conv(x, layer.weight, layer.bias, 1),
    "F.conv2d": lambda: F.conv2d(x, layer.weight, layer.bias, padding=1),
}
with torch.no_grad():
    for name, function in functions.items():
        for _ in range(3):
            function()  # Warm-up.
        timer = benchmark.Timer(stmt="function()", globals={"function": function})
        print(f"{name:<18}: {timer.blocked_autorange(min_run_time=0.5).median * 1e3:6.1f} ms")

Discuss. The view itself costs no memory. Does that mean the einsum never copies? (Try torch.einsum on a non-contiguous input and think about how it is turned into a matrix product.) Why is F.conv2d still the fastest?

Optional C: measuring receptive fields ★★

After Part 1. The receptive field of a unit is the set of input pixels that can affect it. In the lecture we computed it by hand. Here we measure it with autograd: the gradient of one output unit with respect to the input is nonzero exactly on the pixels the unit depends on.

TODO (★★): complete receptive_field(network, size). Run network on a random input of shape (1, 1, size, size) that requires gradients, take the output unit of channel 0 at the center of the output (out[0, 0, h // 2, w // 2]), compute its gradient with respect to the input, and return the height and width of the smallest rectangle containing all its nonzero entries.

To avoid gradients that are zero by chance (after a ReLU or a max-pooling), use networks made of convolutions only, with all-positive weights.

def receptive_field(network, size: int = 64) -> tuple[int, int]:
    # TODO: complete the code.
    raise NotImplementedError("Optional exercise: receptive_field")

def positive_convs(*specs):  # (kernel, stride) for each layer.
    layers = [MyConv2d(1, 1, k, stride=stride, padding=k // 2) for k, stride in specs]
    with torch.no_grad():
        for layer in layers:
            layer.weight.uniform_(0.1, 1.)
    return torch.nn.Sequential(*layers)

assert receptive_field(positive_convs((3, 1))) == (3, 3)
assert receptive_field(positive_convs((5, 1))) == (5, 5)
for specs in [[(3, 1)] * 2, [(3, 1)] * 3, [(3, 1)] * 5, [(3, 2), (3, 1)], [(3, 2), (3, 2), (3, 1)], [(5, 1), (3, 1)]]:
    print(f"{specs}: receptive field {receptive_field(positive_convs(*specs))}")

Discuss. How does the receptive field grow when you stack 3×3 layers with stride 1? What changes after a layer with stride 2? Write a formula, and use it to compute the receptive field of the last convolution of MyCNN (treating each max-pooling as a 2×2 layer with stride 2).

Optional D: causal convolutions ★★

After Part 1, together with the second half of the CNN lecture. Convolutions also work on sequences: a 1D convolution slides a filter of size $k$ along the time axis of an input $(B, C, T)$. For next-token prediction, the output at time $t$ must not see the inputs after $t$, otherwise the model can simply read the answer. A causal convolution pads $k - 1$ zeros on the left only, so that each output looks at the current step and the $k - 1$ previous ones.

TODO (★★): complete causal_conv1d(x, weight, bias), with weight of shape $(C', C, k)$ as in F.conv1d. Use F.pad and the tensor method Tensor.unfold (not the same as F.unfold!) to extract the windows, then an einsum.

def causal_conv1d(x: Float[torch.Tensor, "batch channels time"], weight: Float[torch.Tensor, "out_channels channels kernel"], bias: Float[torch.Tensor, "out_channels"]) -> Float[torch.Tensor, "batch out_channels time"]:
    # TODO: complete the code.
    raise NotImplementedError("Optional exercise: causal_conv1d")

gen = torch.Generator().manual_seed(10)
x, weight, bias = torch.randn(2, 3, 12, generator=gen), torch.randn(4, 3, 5, generator=gen), torch.randn(4, generator=gen)
torch.testing.assert_close(causal_conv1d(x, weight, bias), F.conv1d(F.pad(x, (4, 0)), weight, bias))

# Which inputs does each output depend on? Row t of the Jacobian is output t.
def dependence(function, T=12):
    jac = torch.autograd.functional.jacobian(lambda s: function(s[None, None])[0, 0], torch.randn(T, generator=gen))
    return (jac != 0).float()

w1, b1 = torch.rand(1, 1, 5, generator=gen) + 0.1, torch.zeros(1)
causal = dependence(lambda s: causal_conv1d(s, w1, b1))
centered = dependence(lambda s: F.conv1d(s, w1, b1, padding=2))
assert causal.triu(diagonal=1).sum() == 0, "A causal convolution cannot depend on future inputs"
fig, axes = plt.subplots(1, 2, figsize=(7, 3.4))
for ax, matrix, title in zip(axes, [causal, centered], ["Causal", "Centered (padding 2)"]):
    ax.imshow(matrix, cmap="Blues")
    ax.set(title=title, xlabel="Input time step", ylabel="Output time step")
fig.tight_layout()
plt.show()

Discuss. In the centered convolution, output $t$ depends on inputs $t+1$ and $t+2$. If we trained it to predict the next token, why would the training loss be misleadingly low? How does the receptive field of a stack of causal layers grow, and how could you make it grow faster without more layers?

Optional E: data augmentation and MixUp ★★

After Part 2. The original version of this lab started with a question: which data augmentations make sense for this dataset?

TODO (★): choose a set of suitable augmentations from torchvision's transforms and complete augment below, which receives a batch of standardized images. Think about which transformations keep a character recognizable: a horizontal flip, for instance, would turn many of them into a different character (or into no character at all).

TODO (★★): complete mixup(x, y, alpha). Sample $\lambda \sim \text{Beta}(\alpha, \alpha)$ once per batch and a random permutation of the batch, and return the mixed images $\lambda x + (1 - \lambda) x_\pi$ together with the mixed one-hot targets $\lambda\, \text{onehot}(y) + (1 - \lambda)\, \text{onehot}(y_\pi)$. F.cross_entropy accepts these soft targets directly.

The cell after the checks trains the same residual network with no regularization, with your augmentations, and with MixUp, on 10,000 training images.

def augment(x: Float[torch.Tensor, "batch 1 height width"]) -> Float[torch.Tensor, "batch 1 height width"]:
    # TODO: complete the code.
    raise NotImplementedError("Optional exercise: augment")

def mixup(x: Float[torch.Tensor, "batch channels height width"], y: Int[torch.Tensor, "batch"], alpha: float = 0.4, n_classes: int = 10) -> tuple[torch.Tensor, torch.Tensor]:
    # TODO: complete the code.
    raise NotImplementedError("Optional exercise: mixup")

torch.manual_seed(11)
augmented = augment(xb[:8])
assert augmented.shape == xb[:8].shape and not torch.equal(augmented, xb[:8]), "augment should change the images"
mixed_x, mixed_y = mixup(xb, yb)
assert mixed_x.shape == xb.shape and mixed_y.shape == (64, 10)
torch.testing.assert_close(mixed_y.sum(dim=1), torch.ones(64), msg="Mixed targets must sum to 1")
pairs = mixed_y[(mixed_y > 0).sum(dim=1) == 2].sort(dim=1, descending=True).values[:, :2]
assert torch.allclose(pairs, pairs[:1].expand_as(pairs)), "Use a single lambda for the whole batch"
fig, axes = plt.subplots(1, 8, figsize=(8, 1.3))
for ax, image in zip(axes, augmented):
    ax.imshow(image[0], cmap="gray_r")
    ax.axis("off")
plt.show()

class Regularized(torch.nn.Module):
    # Augmentation as a layer of the model, active only in training mode (as in the lecture).
    def __init__(self, network, mode):
        super().__init__()
        self.network, self.mode = network, mode

    def forward(self, x):
        if self.training and self.mode == "augment":
            x = augment(x.cpu()).to(x.device)
        return self.network(x)

def fit_mixup(model, loader, epochs, lr=1e-3):
    # The same loop as fit, with mixed inputs and soft targets.
    model.to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
    history = {"train_loss": [], "val_loss": [], "val_accuracy": []}
    for _ in range(epochs):
        model.train()
        for x, y in loader:
            x, soft = mixup(x.to(device), y.to(device))
            optimizer.zero_grad()
            loss = F.cross_entropy(model(x), soft)
            loss.backward()
            optimizer.step()
            history["train_loss"].append(loss.item())
        val_loss, val_accuracy = evaluate(model, val_loader)
        history["val_loss"].append(val_loss)
        history["val_accuracy"].append(val_accuracy)
    model.cpu()
    return history

REGULARIZATION_EPOCHS = 6
regularization = {}
for mode in ["none", "augment", "mixup"]:
    torch.manual_seed(0)
    model = Regularized(deep_cnn(3, residual=True, batchnorm=True, channels=32), mode)
    start = time.time()
    if mode == "mixup":
        regularization[mode] = fit_mixup(model, subset_loader, REGULARIZATION_EPOCHS)
    else:
        regularization[mode] = fit(model, loader=subset_loader, epochs=REGULARIZATION_EPOCHS, verbose=False)
    print(f"{mode:<8}: validation loss {regularization[mode]['val_loss'][-1]:.3f}, "
          f"accuracy {regularization[mode]['val_accuracy'][-1]:.3f} ({time.time() - start:.0f} s)")

fig, axes = plt.subplots(1, 2, figsize=(10, 3.5))
for mode, history in regularization.items():
    epochs = range(1, len(history["val_loss"]) + 1)
    axes[0].plot(epochs, history["val_loss"], "o-", label=mode)
    axes[1].plot(epochs, history["val_accuracy"], "o-", label=mode)
axes[0].set(xlabel="Epoch", ylabel="Validation loss")
axes[1].set(xlabel="Epoch", ylabel="Validation accuracy")
axes[0].legend()
fig.tight_layout()
plt.show()

Discuss. Which variant has the lowest validation loss, and which the highest accuracy? With MixUp, the training loss is never close to zero, even for a perfect classifier. Why? Do you expect the ranking to change with more epochs, or with all 54,000 training images?

Sources

• Simone Scardapane, original Colab lab: Convolutional neural networks (from scratch).
• NNDS 2026 lectures on convolutional neural networks and on deeper models (normalization, residual connections, regularization).
• PT01 (views and strides, autograd) and PT02 (the training loop).
• KMNIST: Clanuwat et al. (2018), Deep learning for classical Japanese literature, dataset at github.com/rois-codh/kmnist (CC BY-SA 4.0).
• Occlusion maps: Zeiler and Fergus (2014), Visualizing and understanding convolutional networks, ECCV.
• Batch normalization: Ioffe and Szegedy (2015), ICML. Residual networks: He et al. (2016), CVPR. Dropout: Srivastava et al. (2014), JMLR. MixUp: Zhang et al. (2018), ICLR.
• PyTorch documentation for F.unfold, Tensor.as_strided and torch.utils.benchmark.