Convolutional Networks (CNN)

Code

Pooling and the CNN — write them yourself

The convolution works. Two parts remain — this time with no holes, just empty skeletons: pooling and the CNN itself.

You have already seen both in step 2.

Pooling

out ← table of zeros [H // size, W // size]
for each block (i, j):
    out[i, j] ← the largest number in the block
               x[i*size .. (i+1)*size, j*size .. (j+1)*size]
return out

Blocks do not overlap and move in steps of size (the convolution moved by 1). If the width is not divisible by size, the last incomplete column or row is dropped — nn.MaxPool2d does the same. The test checks it with a 5 × 5 input.

`CNN`

This is the same stack you saw in the shape table:

Conv2d(1, c1, 3, padding=1) → ReLU → MaxPool2d(2)      28 → 14
Conv2d(c1, c2, 3, padding=1) → ReLU → MaxPool2d(2)     14 → 7
Flatten → Dropout → Linear(c2 · 7 · 7, num_classes)

Three places where people go wrong (and what the test will tell you):

The Linear input. After two poolings the area is 7 × 7, not 14 × 14. Write c2 * 14 * 14 and you get mat1 and mat2 shapes cannot be multiplied.
nn.Flatten(). Without it Linear receives a 4-dimensional tensor [N, 32, 7, 7] instead of a row [N, 1568].
The start of forward. All your code so far, and fashion.py, supplies flat [N, 784] rows, while Conv2d wants [N, 1, 28, 28]. One reshape line at the start — and CNN accepts both shapes.

The parameter count is a check for yourself: 160 (first conv) + 4,640 (second)

15,690 (dense) = 20,490.

Your job

Two files grew, one new skeleton in each:

conv_by_hand.py gained max_pool2d_by_hand with raise NotImplementedError;
model.py gained the class CNN — with an __init__ and a forward that do not work yet. Note: the docstring lists everything the test demands.

conv_by_hand.py is still shown with its holes in the panel — the panel shows the skeleton, not your file. In your copy they are already filled; leave it that way.

test_cnn.py grew too: it now checks pooling (its numbers come from step 2: the ReLU map → [[36, 0], [27, 0]]) and CNN (the parameter count, two each of Conv2d, ReLU and MaxPool2d, that it accepts [N, 784] and [N, 1, 28, 28], and that it learns: in 60 steps on 16 random images the loss must fall by more than a factor of 4).

Run python -m unittest test_cnn. To begin with you will see:

EEEEEE...EEEE
======================================================================
ERROR: test_accepts_images_and_flat_rows (test_cnn.CNNTest.test_accepts_images_and_flat_rows)
NotImplementedError: build the layers

(...)
ERROR: test_a_leftover_row_and_column_are_dropped (test_cnn.MaxPoolByHandTest.test_a_leftover_row_and_column_are_dropped)
NotImplementedError: write max_pool2d_by_hand

(...)
Ran 13 tests in 0.6s

FAILED (errors=10)

The three . are the three convolution tests you already passed. The ten NotImplementedErrors are what is left to write. (Output shortened.)

When it does not pass

mat1 and mat2 shapes cannot be multiplied (4x1568 and 6272x10) — the Linear was built with c2 * 14 * 14 (6,272), but the network delivers c2 * 7 * 7 (1,568).
mat1 and mat2 shapes cannot be multiplied (3584x7 and 1568x10) — you forgot nn.Flatten().
Expected 3D (unbatched) or 4D (batched) input to conv2d, but got input of size: [4, 784] — you forgot the reshape at the start of forward.
AssertionError: 0 != 2 in test_has_the_layers… — you left out the ReLUs. Without them two convolutions in a row are one (lesson 6).
Pooling test: Tensor-likes are not close or a wrong shape — blocks taken with stride 1 instead of size, or an average used instead of .max().

conv_by_hand.py MODIFY 10 lines changed

Edit this file in ./

"""conv_by_hand.py — a convolution written with plain loops."""

import torch


def conv2d_by_hand(image: torch.Tensor, kernel: torch.Tensor) -> torch.Tensor:
    """Slide kernel over image (no padding, stride 1) and return the feature map.

    image: [H, W]   kernel: [kh, kw]   result: [H - kh + 1, W - kw + 1]
    Every output number is the sum of (patch * kernel), where patch is the
    kh x kw piece of the image that lies under the kernel at that position.
    """
    h, w = image.shape
    kh, kw = kernel.shape
    out = torch.zeros(h - kh + 1, w - kw + 1)
    for i in range(out.shape[0]):
        for j in range(out.shape[1]):
            # HOLE 1: cut out the kh x kw patch whose top-left corner is (i, j).
            patch = image[:kh, :kw]
            # HOLE 2: multiply patch and kernel element by element, add it all up.
            out[i, j] = 0.0
    return out


def max_pool2d_by_hand(x: torch.Tensor, size: int = 2) -> torch.Tensor:
    """Keep the largest number in every size x size block (stride = size).

    x: [H, W]   result: [H // size, W // size]
    A leftover row or column that does not fill a whole block is dropped,
    exactly as nn.MaxPool2d does it.
    """
    raise NotImplementedError("write max_pool2d_by_hand")

model.py MODIFY 24 lines changed

Edit this file in ./

"""model.py — the model as an nn.Module."""

import torch
from torch import nn


class LinearClassifier(nn.Module):
    """The same 10 neurons as lesson 4: scores = x @ W + b."""

    def __init__(self, in_features: int = 784, num_classes: int = 10) -> None:
        super().__init__()
        self.fc = nn.Linear(in_features, num_classes)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.fc(x)  # raw scores (logits); the loss applies softmax itself


class MLP(nn.Module):
    """Linear -> ReLU -> (Dropout) -> ... -> Linear."""

    def __init__(
        self,
        in_features: int = 784,
        hidden: tuple[int, ...] = (128,),
        num_classes: int = 10,
        dropout: float = 0.0,
    ) -> None:
        super().__init__()
        layers: list[nn.Module] = []
        width = in_features
        for h in hidden:
            layers += [nn.Linear(width, h), nn.ReLU(), nn.Dropout(dropout)]
            width = h
        layers.append(nn.Linear(width, num_classes))
        self.net = nn.Sequential(*layers)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


class CNN(nn.Module):
    """conv -> ReLU -> pool, twice, then one Linear layer on the feature maps.

    Input: [N, 1, 28, 28] images, or flat [N, 784] rows (forward reshapes them).
    Layers, in order:
      Conv2d(1, c1, kernel_size=3, padding=1) -> ReLU -> MaxPool2d(2)    28 -> 14
      Conv2d(c1, c2, kernel_size=3, padding=1) -> ReLU -> MaxPool2d(2)   14 -> 7
      Flatten -> Dropout(dropout) -> Linear(c2 * 7 * 7, num_classes)
    With the defaults (16, 32) that is 20,490 parameters.
    """

    def __init__(
        self,
        num_classes: int = 10,
        channels: tuple[int, int] = (16, 32),
        dropout: float = 0.0,
    ) -> None:
        super().__init__()
        raise NotImplementedError("build the layers")

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        raise NotImplementedError("reshape to [N, 1, 28, 28], then run the layers")


def count_parameters(model: nn.Module) -> int:
    return sum(p.numel() for p in model.parameters())

test_cnn.py MODIFY 67 lines changed

Edit this file in ./

"""test_cnn.py — the verdict. Run: python -m unittest -v"""

import unittest

import torch
import torch.nn.functional as F

from conv_by_hand import conv2d_by_hand, max_pool2d_by_hand
from model import CNN, count_parameters

# The 6x6 picture from the lesson: a bright 3x3 square on a dark background.
SQUARE = torch.tensor(
    [
        [0, 0, 0, 0, 0, 0],
        [0, 9, 9, 9, 0, 0],
        [0, 9, 9, 9, 0, 0],
        [0, 9, 9, 9, 0, 0],
        [0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0],
    ],
    dtype=torch.float32,
)
# Responds where brightness grows from left to right (a vertical edge).
EDGE = torch.tensor([[-1.0, 0.0, 1.0], [-2.0, 0.0, 2.0], [-1.0, 0.0, 1.0]])


class ConvByHandTest(unittest.TestCase):
    def test_output_shape(self):
        out = conv2d_by_hand(torch.zeros(28, 28), EDGE)
        self.assertEqual(tuple(out.shape), (26, 26))

    def test_the_numbers_from_the_lesson(self):
        out = conv2d_by_hand(SQUARE, EDGE)
        expected = torch.tensor(
            [
                [27.0, 0.0, -27.0, -27.0],
                [36.0, 0.0, -36.0, -36.0],
                [27.0, 0.0, -27.0, -27.0],
                [9.0, 0.0, -9.0, -9.0],
            ]
        )
        self.assertTrue(torch.equal(out, expected), f"got\n{out}")

    def test_matches_torch_on_random_input(self):
        g = torch.Generator().manual_seed(0)
        image = torch.randn(9, 11, generator=g)
        kernel = torch.randn(3, 3, generator=g)
        want = F.conv2d(image.reshape(1, 1, 9, 11), kernel.reshape(1, 1, 3, 3))[0, 0]
        torch.testing.assert_close(conv2d_by_hand(image, kernel), want)


class MaxPoolByHandTest(unittest.TestCase):
    def test_the_numbers_from_the_lesson(self):
        fmap = F.relu(conv2d_by_hand(SQUARE, EDGE))
        out = max_pool2d_by_hand(fmap)
        expected = torch.tensor([[36.0, 0.0], [27.0, 0.0]])
        self.assertTrue(torch.equal(out, expected), f"got\n{out}")

    def test_halves_each_side(self):
        out = max_pool2d_by_hand(torch.zeros(26, 26))
        self.assertEqual(tuple(out.shape), (13, 13))

    def test_a_leftover_row_and_column_are_dropped(self):
        out = max_pool2d_by_hand(torch.arange(25.0).reshape(5, 5))
        self.assertEqual(tuple(out.shape), (2, 2))

    def test_matches_torch_on_random_input(self):
        x = torch.randn(8, 8, generator=torch.Generator().manual_seed(1))
        want = F.max_pool2d(x.reshape(1, 1, 8, 8), 2)[0, 0]
        torch.testing.assert_close(max_pool2d_by_hand(x), want)


class CNNTest(unittest.TestCase):
    def test_default_parameter_count(self):
        # 160 + 4,640 + 15,690 = 20,490
        self.assertEqual(count_parameters(CNN()), 20_490)

    def test_has_the_layers_the_contract_names(self):
        # Without the ReLUs the whole stack collapses into one linear map (lesson 6).
        kinds = [type(m) for m in CNN().modules()]
        self.assertEqual(kinds.count(torch.nn.Conv2d), 2)
        self.assertEqual(kinds.count(torch.nn.ReLU), 2)
        self.assertEqual(kinds.count(torch.nn.MaxPool2d), 2)
        self.assertEqual(kinds.count(torch.nn.Linear), 1)

    def test_wider_network_has_more_parameters(self):
        self.assertEqual(count_parameters(CNN(channels=(32, 64))), 50_186)

    def test_accepts_images_and_flat_rows(self):
        model = CNN().eval()
        images = torch.rand(4, 1, 28, 28)
        out_images = model(images)
        out_flat = model(images.reshape(4, 784))
        self.assertEqual(tuple(out_images.shape), (4, 10))
        torch.testing.assert_close(out_flat, out_images)

    def test_output_width_follows_num_classes(self):
        out = CNN(num_classes=4)(torch.rand(2, 1, 28, 28))
        self.assertEqual(tuple(out.shape), (2, 4))

    def test_can_memorise_a_tiny_batch(self):
        torch.manual_seed(0)
        x = torch.rand(16, 784)
        y = torch.randint(0, 10, (16,))
        model = CNN()
        optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
        loss_fn = torch.nn.CrossEntropyLoss()
        first = loss_fn(model(x), y).item()
        for _ in range(60):
            loss = loss_fn(model(x), y)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        self.assertLess(loss.item(), first / 4, f"loss {first:.3f} -> {loss.item():.3f}")


if __name__ == "__main__":
    unittest.main()