%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__)

PyTorch 2.14.1 · torchvision 0.29.1

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()

60000 training images, 10000 test images
<class 'PIL.Image.Image'> 8 re

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.

Pixel mean 0.1918, std 0.3483 (training images)
torch.Size([64, 1, 28, 28]) torch.Size([64])

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()

Training on cpu

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()))

torch.Size([64, 16, 14, 14])
weight (16, 1, 5, 5)
bias (16,)
Parameters: 416

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)

MyLinear()
[('weight', (16, 32)), ('bias', (32,))]

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

tensor([[[[1., 2., 3.],
          [4., 5., 6.],
          [7., 8., 9.]]]])
torch.Size([1, 4, 4])
tensor([[[1., 2., 4., 5.],
         [2., 3., 5., 6.],
         [4., 5., 7., 8.],
         [5., 6., 8., 9.]]])

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

torch.Size([64, 9, 784])

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")

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__()
        self.kernel_size, self.stride, self.padding = kernel_size, stride, padding
        fan_in = in_channels * kernel_size ** 2
        self.weight = Parameter(torch.randn(out_channels, in_channels, kernel_size, kernel_size) * math.sqrt(2 / fan_in))
        self.bias = Parameter(torch.zeros(out_channels))

    def forward(self, xb: Float[torch.Tensor, "batch in_channels height width"]) -> Float[torch.Tensor, "batch out_channels new_height new_width"]:
        B, _, H, W = xb.shape
        s, p, stride = self.kernel_size, self.padding, self.stride
        new_height, new_width = (H + 2 * p - s) // stride + 1, (W + 2 * p - s) // stride + 1
        U = F.unfold(xb, s, padding=p, stride=stride)          # (B, C s^2, L)
        Z = self.weight.view(self.weight.shape[0], -1) @ U    # (B, C', L)
        Z = Z + self.bias[:, None]
        return Z.view(B, -1, new_height, new_width)

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.")

Activity 1 checks passed.

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")

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"]:
        B, C, H, W = xb.shape
        k = self.kernel_size
        patches = F.unfold(xb, k, stride=k).view(B, C, k * k, -1)
        return patches.max(dim=2).values.view(B, C, H // k, W // k)

def global_avg_pool(xb: Float[torch.Tensor, "batch channels height width"]) -> Float[torch.Tensor, "batch channels"]:
    return xb.mean(dim=(2, 3))

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.")

Activity 2 checks passed.

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())

tensor([[-0.077,  0.360, -0.782,  0.072],
        [ 0.665, -0.287,  1.621, -1.597],
        [-0.052, -0.306,  0.249, -0.223],
        [ 0.913,  0.204,  0.574,  0.416]])
tensor([[0., 0., 0., 0.],
        [1., 0., 1., 0.],
        [0., 0., 0., 0.],
        [1., 0., 1., 0.]])

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")

class MyDepthwiseConv2d(torch.nn.Module):
    def __init__(self, channels: int, kernel_size: int, padding: int = 0):
        super().__init__()
        self.kernel_size, self.padding = kernel_size, padding
        self.weight = Parameter(torch.randn(channels, 1, kernel_size, kernel_size) * math.sqrt(2 / kernel_size ** 2))
        self.bias = Parameter(torch.zeros(channels))

    def forward(self, xb: Float[torch.Tensor, "batch channels height width"]) -> Float[torch.Tensor, "batch channels new_height new_width"]:
        B, C, H, W = xb.shape
        s, p = self.kernel_size, self.padding
        patches = F.unfold(xb, s, padding=p).view(B, C, s * s, -1)
        Z = (patches * self.weight.view(C, s * s, 1)).sum(dim=2) + self.bias[:, None]
        return Z.view(B, C, H + 2 * p - s + 1, W + 2 * p - s + 1)

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.")

Activity 3 checks passed.

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)

Standard:  73856
Separable: 8960

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")

class MyCNN(torch.nn.Module):
    def __init__(self, n_classes: int = 10):
        super().__init__()
        self.conv1 = MyConv2d(1, 32, 3, padding=1)
        self.conv2 = MyConv2d(32, 64, 3, padding=1)
        self.conv3 = MyConv2d(64, 64, 3, padding=1)
        self.pool = MyMaxPool2d(2)
        self.head = MyLinear(64, n_classes)

    def forward(self, xb: Float[torch.Tensor, "batch channels height width"]) -> Float[torch.Tensor, "batch n_classes"]:
        h = self.pool(F.relu(self.conv1(xb)))   # (B, 32, 14, 14)
        h = self.pool(F.relu(self.conv2(h)))    # (B, 64, 7, 7)
        h = F.relu(self.conv3(h))               # (B, 64, 7, 7)
        return self.head(global_avg_pool(h))

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.")

Activity 4 checks passed: 56,394 parameters.

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

Epoch 1: validation loss 0.228, accuracy 0.937 (26 s)

Epoch 2: validation loss 0.152, accuracy 0.955 (26 s)

Epoch 3: validation loss 0.121, accuracy 0.965 (26 s)

Epoch 4: validation loss 0.105, accuracy 0.968 (26 s)

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}")

Epoch 1: validation loss 0.235, accuracy 0.930 (3 s)

Epoch 2: validation loss 0.221, accuracy 0.937 (3 s)

Epoch 3: validation loss 0.211, accuracy 0.940 (3 s)

Epoch 4: validation loss 0.224, accuracy 0.942 (3 s)

CNN:  56,394 parameters, validation accuracy 0.968

MLP: 101,770 parameters, validation accuracy 0.942

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()

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")

@torch.no_grad()
def equivariance_error(layer, x: Float[torch.Tensor, "batch channels height width"], dy: int, dx: int, border: int) -> float:
    difference = layer(shift(x, dy, dx)) - shift(layer(x), dy, dx)
    if border > 0:
        difference = difference[..., border:-border, border:-border]
    return difference.abs().max().item()

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.")

Activity 5 checks passed.

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

for border in range(5):
    print("One layer, border", border, ":", equivariance_error(cnn.conv1, images, 2, 2, border))
two_layers = lambda x: cnn.conv2(F.relu(cnn.conv1(x)))
for border in range(5):
    print("Two layers, border", border, ":", equivariance_error(two_layers, images, 2, 2, border))
pool = MyMaxPool2d(2)
x0 = torch.randn(8, 1, 28, 28)  # Random values, so that ties between pixels are unlikely.
for d in [2, 1]:
    error = (pool(shift(x0, d, d)) - shift(pool(x0), d // 2, d // 2))[..., 2:-2, 2:-2].abs().max().item()
    print(f"Max-pooling, shift {d}: error {error:.3f}")

One layer, border 0 : 6.369101047515869
One layer, border 1 : 4.775282859802246
One layer, border 2 : 0.0
One layer, border 3 : 0.0
One layer, border 4 : 0.0
Two layers, border 0 : 64.28025817871094
Two layers, border 1 : 36.975826263427734
Two layers, border 2 : 15.55029296875
Two layers, border 3 : 0.0
Two layers, border 4 : 0.0
Max-pooling, shift 2: error 0.000
Max-pooling, shift 1: error 3.576

@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()

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()

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")

@torch.no_grad()
def occlusion_map(model, image: Float[torch.Tensor, "1 height width"], label: int, patch: int = 7, stride: int = 3) -> Float[torch.Tensor, "rows cols"]:
    model.eval()
    H, W = image.shape[-2:]
    rows, cols = range(0, H - patch + 1, stride), range(0, W - patch + 1, stride)
    batch = image.repeat(len(rows) * len(cols), 1, 1, 1)
    for n, (i, j) in enumerate((i, j) for i in rows for j in cols):
        batch[n, :, i:i + patch, j:j + patch] = background
    probabilities = model(batch).softmax(dim=1)[:, label]
    return probabilities.view(len(rows), len(cols))

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.")

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()

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}")

CNN: validation accuracy 0.968, test accuracy 0.916

MLP: validation accuracy 0.942, test accuracy 0.871

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")

class MyBatchNorm2d(torch.nn.Module):
    def __init__(self, channels: int, eps: float = 1e-5, momentum: float = 0.1):
        super().__init__()
        self.eps, self.momentum = eps, momentum
        self.weight = Parameter(torch.ones(channels))
        self.bias = Parameter(torch.zeros(channels))
        self.register_buffer("running_mean", torch.zeros(channels))
        self.register_buffer("running_var", torch.ones(channels))

    def forward(self, xb: Float[torch.Tensor, "batch channels height width"]) -> Float[torch.Tensor, "batch channels height width"]:
        if self.training:
            mean = xb.mean(dim=(0, 2, 3))
            var = xb.var(dim=(0, 2, 3), correction=0)
            with torch.no_grad():
                n = xb.numel() / xb.shape[1]
                self.running_mean.lerp_(mean, self.momentum)
                self.running_var.lerp_(var * n / (n - 1), self.momentum)
        else:
            mean, var = self.running_mean, self.running_var
        normalized = (xb - mean[:, None, None]) / torch.sqrt(var[:, None, None] + self.eps)
        return self.weight[:, None, None] * normalized + self.bias[:, None, None]

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.")

Activity 7 checks passed.

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

def fuse_conv_bn(conv: MyConv2d, bn: MyBatchNorm2d) -> MyConv2d:
    fused = copy.deepcopy(conv)
    with torch.no_grad():
        a = bn.weight / torch.sqrt(bn.running_var + bn.eps)
        fused.weight.copy_(conv.weight * a[:, None, None, None])
        fused.bias.copy_(a * conv.bias + bn.bias - a * bn.running_mean)
    return fused

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.")

Activity 8 checks passed.

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)

Epoch 1: validation loss 0.289, accuracy 0.917 (32 s)

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())

Alone:                     0.07519209384918213
With 7 of the same class:  0.40791055560112
With 7 of other classes:   0.9832617044448853

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])

Validation accuracy before: 0.9175

Validation accuracy after:  0.25566666666666665

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")

@torch.no_grad()
def safe_evaluate(model, x: Float[torch.Tensor, "batch channels height width"]) -> Float[torch.Tensor, "batch n_classes"]:
    was_training = model.training
    before = {name: buffer.clone() for name, buffer in model.named_buffers()}
    model.eval()
    try:
        probabilities = model(x).softmax(dim=1)
    finally:
        model.train(was_training)
    for name, buffer in model.named_buffers():
        assert torch.equal(buffer, before[name]), f"Evaluation changed the buffer {name}"
    return probabilities

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.")

Activity 9 checks passed.

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")

class ResidualBlock(torch.nn.Module):
    def __init__(self, channels: int, residual: bool = True, batchnorm: bool = True):
        super().__init__()
        self.residual = residual
        norm = (lambda: torch.nn.BatchNorm2d(channels)) if batchnorm else torch.nn.Identity
        self.conv1 = torch.nn.Conv2d(channels, channels, 3, padding=1, bias=not batchnorm)
        self.bn1 = norm()
        self.conv2 = torch.nn.Conv2d(channels, channels, 3, padding=1, bias=not batchnorm)
        self.bn2 = norm()

    def forward(self, xb: Float[torch.Tensor, "batch channels height width"]) -> Float[torch.Tensor, "batch channels height width"]:
        h = self.bn2(self.conv2(F.relu(self.bn1(self.conv1(xb)))))
        return F.relu(xb + h if self.residual else h)

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.")

Activity 10 checks passed.

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()

 6 layers, Plain        : training loss 1.639, validation accuracy 0.431 (8 s)

 6 layers, Plain + BN   : training loss 1.264, validation accuracy 0.721 (9 s)

 6 layers, Residual + BN: training loss 1.213, validation accuracy 0.669 (9 s)

18 layers, Plain        : training loss 2.301, validation accuracy 0.102 (27 s)

18 layers, Plain + BN   : training loss 0.686, validation accuracy 0.830 (29 s)

18 layers, Residual + BN: training loss 0.641, validation accuracy 0.788 (29 s)

34 layers, Plain        : training loss 2.303, validation accuracy 0.101 (52 s)

34 layers, Plain + BN   : training loss 1.807, validation accuracy 0.363 (57 s)

34 layers, Residual + BN: training loss 0.780, validation accuracy 0.754 (57 s)

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")

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, "..."]:
        if not self.training or self.p == 0:
            return xb
        mask = (torch.rand_like(xb) > self.p).to(xb.dtype)
        return xb * mask / (1 - self.p)

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.")

Activity 11 checks passed.

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")

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__()
        self.kernel_size, self.padding = kernel_size, padding
        self.new_height, self.new_width = height + 2 * padding - kernel_size + 1, width + 2 * padding - kernel_size + 1
        windows, fan_in = self.new_height * self.new_width, in_channels * kernel_size ** 2
        self.weight = Parameter(torch.randn(windows, out_channels, fan_in) * math.sqrt(2 / fan_in))
        self.bias = Parameter(torch.zeros(windows, out_channels))

    def forward(self, xb: Float[torch.Tensor, "batch in_channels height width"]) -> Float[torch.Tensor, "batch out_channels new_height new_width"]:
        U = F.unfold(xb, self.kernel_size, padding=self.padding)          # (B, C s^2, L)
        Z = torch.einsum("lod,bdl->bol", self.weight, U) + self.bias.T    # (B, C', L)
        return Z.view(xb.shape[0], -1, self.new_height, self.new_width)

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))

Parameters, convolution:        160
Parameters, locally connected:  125440
Equivariance error (border 3):  15.850765228271484

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")

def patches_view(x: Float[torch.Tensor, "batch channels height width"], kernel_size: int) -> Float[torch.Tensor, "batch channels rows cols kernel kernel"]:
    B, C, H, W = x.shape
    sB, sC, sH, sW = x.stride()
    s = kernel_size
    return x.as_strided((B, C, H - s + 1, W - s + 1, s, s), (sB, sC, sH, sW, sH, sW), x.storage_offset())

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())

Is the view contiguous? False

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")

Input: 6.1 MiB, unfold: 55.1 MiB

MyConv2d (unfold) :   20.5 ms

einsum on a view  :   28.8 ms

F.conv2d          :   15.2 ms

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

def receptive_field(network, size: int = 64) -> tuple[int, int]:
    x = torch.randn(1, 1, size, size, requires_grad=True)
    out = network(x)
    h, w = out.shape[-2:]
    out[0, 0, h // 2, w // 2].backward()
    rows, cols = (x.grad[0, 0] != 0).nonzero(as_tuple=True)
    return rows.max().item() - rows.min().item() + 1, cols.max().item() - cols.min().item() + 1

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))}")

[(3, 1), (3, 1)]: receptive field (5, 5)
[(3, 1), (3, 1), (3, 1)]: receptive field (7, 7)
[(3, 1), (3, 1), (3, 1), (3, 1), (3, 1)]: receptive field (11, 11)
[(3, 2), (3, 1)]: receptive field (7, 7)
[(3, 2), (3, 2), (3, 1)]: receptive field (15, 15)
[(5, 1), (3, 1)]: receptive field (7, 7)

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")

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"]:
    k = weight.shape[-1]
    windows = F.pad(x, (k - 1, 0)).unfold(-1, k, 1)    # (B, C, T, k)
    return torch.einsum("bctk,ock->bot", windows, weight) + bias[:, None]

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()

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")

augmentations = v2.RandomAffine(degrees=10, translate=(0.1, 0.1), scale=(0.9, 1.1), fill=background)

def augment(x: Float[torch.Tensor, "batch 1 height width"]) -> Float[torch.Tensor, "batch 1 height width"]:
    return torch.stack([augmentations(image) for image in x])

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]:
    lam = torch.distributions.Beta(alpha, alpha).sample()
    perm = torch.randperm(len(x))
    targets = F.one_hot(y, n_classes).float()
    return lam * x + (1 - lam) * x[perm], lam * targets + (1 - lam) * targets[perm]

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()

none    : validation loss 0.560, accuracy 0.812 (70 s)

augment : validation loss 0.214, accuracy 0.939 (74 s)

mixup   : validation loss 0.316, accuracy 0.923 (70 s)

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	★

PT04 · Convolutional neural networks from scratch (instructor version)¶

1. Setup and data¶

2. Layers are modules¶

Part 1 · Convolutional networks¶

3. Activity 1: a convolutional layer ★★¶

4. Activity 2: pooling ★¶

5. Activity 3: depthwise convolutions ★★¶

6. Activity 4: building the network ★¶

7. Activity 5: translation equivariance ★★¶

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

9. The test set¶

Part 2 · Training deeper networks¶

10. Activity 7: batch normalization ★★¶

11. Activity 8: fusing batch normalization ★★¶

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

13. Activity 10: residual blocks ★¶

14. Activity 11: dropout ★¶

What we built¶

Optional A: a locally-connected layer ★★★¶

Optional B: patches without copies ★★★¶

Optional C: measuring receptive fields ★★¶

Optional D: causal convolutions ★★¶

Optional E: data augmentation and MixUp ★★¶

Sources¶