Lesson 7: Computer vision basics — convolutional networks (CNN)

import torch
import torch.nn as nn

class SimpleCNN(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()

        # --- spatial feature extractor ---
        self.feature_extractor = nn.Sequential(
            # Conv 1: 1 input channel (greyscale) -> 16 feature channels.
            # padding=1 keeps the size, so the output is 16 x 28 x 28.
            nn.Conv2d(in_channels=1, out_channels=16, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),   # halves it -> 16 x 14 x 14

            # Conv 2: 16 channels -> 32 channels, output 32 x 14 x 14
            nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)    # halves it again -> 32 x 7 x 7
        )

        # --- linear classifier ---
        # After extraction the tensor is (batch, 32, 7, 7); flattening gives
        # (batch, 32*7*7 = 1568). That 1568 is where the shape error comes from.
        self.classifier = nn.Sequential(
            nn.Linear(32 * 7 * 7, 128),   # must match the flattened size exactly
            nn.ReLU(),
            nn.Linear(128, num_classes)   # raw logits, one per class
        )

    def forward(self, x):
        features = self.feature_extractor(x)
        # start_dim=1 keeps the batch dimension and flattens everything after it.
        flat_features = torch.flatten(features, start_dim=1)
        return self.classifier(flat_features)

# Parameters in the two Conv2d layers defined above.
# Conv1: 16 filters, each (1 input channel * 3 * 3) weights + 1 bias
conv1_params = 16 * (1 * 3 * 3 + 1)   # = 160
# Conv2: 32 filters, each (16 input channels * 3 * 3) weights + 1 bias
conv2_params = 32 * (16 * 3 * 3 + 1)  # = 4,640
feature_extractor_total = conv1_params + conv2_params
print(f"two Conv2d layers (the feature extractor): {feature_extractor_total:,}")  # 4,800

# For comparison: the SIMPLEST possible linear layer taking the same image
# already flattened (28*28 = 784) to 128 hidden neurons — matching the width of
# the classifier behind it, and extracting no spatial structure whatsoever.
mlp_equivalent_params = 784 * 128 + 128
print(f"one equivalent linear layer (784 -> 128):  {mlp_equivalent_params:,}")  # 100,480

print(f"the CNN uses {mlp_equivalent_params / feature_extractor_total:.1f}x fewer")  # ~21x

# train_mnist_cnn.py
# Lesson 7: Computer vision basics — convolutional networks
# Practical AI Engineer series
#
# Run it with:  python train_mnist_cnn.py
# Requires:     pip install torch numpy       (no download, no torchvision)
#
# WHY THE DATA IS SYNTHETIC, AND WHAT THAT COSTS
# Real MNIST needs a ~10 MB download and torchvision, which makes the lesson fail
# on a bad connection. So this script draws its own dataset: ten distinct shapes
# standing in for ten digit classes, plus noise and a random offset.
#
# The important part: these shapes are GENUINELY LEARNABLE. An earlier version of
# this script fed the network `np.random.randn` images with `np.random.randint`
# labels — pure noise with no relationship between image and label. That task is
# unlearnable by construction, so accuracy sat at chance level while the training
# loss still fell, because the network was memorising 200 random labels. A falling
# loss with chance-level accuracy is the signature of exactly that mistake.
#
# To train on real MNIST instead, install torchvision and replace
# make_shape_dataset() with torchvision.datasets.MNIST. Everything else is unchanged.

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim

CLASS_NAMES = [
    'ring',
    'vertical bar',
    'horizontal bar',
    'plus',
    'diagonal \\',
    'diagonal /',
    'cross X',
    'hollow box',
    'solid box',
    'double bar',
]


def draw_shape(cls, rng):
    """One 28x28 shape for the given class, jittered slightly off centre."""
    img = np.zeros((28, 28), dtype=np.float32)
    cy, cx = rng.integers(11, 17), rng.integers(11, 17)
    t = 2  # stroke thickness
    clip = lambda v: int(np.clip(v, 0, 27))

    if cls == 0:  # ring
        yy, xx = np.ogrid[:28, :28]
        r = np.sqrt((yy - cy) ** 2 + (xx - cx) ** 2)
        img[(r > 6) & (r < 6 + t + 1)] = 1
    elif cls == 1:  # vertical bar
        img[cy - 9 : cy + 9, cx - 1 : cx + t] = 1
    elif cls == 2:  # horizontal bar
        img[cy - 1 : cy + t, cx - 9 : cx + 9] = 1
    elif cls == 3:  # plus
        img[cy - 9 : cy + 9, cx - 1 : cx + t] = 1
        img[cy - 1 : cy + t, cx - 9 : cx + 9] = 1
    elif cls == 4:  # diagonal \
        for k in range(-9, 9):
            img[clip(cy + k), clip(cx + k)] = 1
    elif cls == 5:  # diagonal /
        for k in range(-9, 9):
            img[clip(cy + k), clip(cx - k)] = 1
    elif cls == 6:  # cross X
        for k in range(-9, 9):
            img[clip(cy + k), clip(cx + k)] = 1
            img[clip(cy + k), clip(cx - k)] = 1
    elif cls == 7:  # hollow box
        img[cy - 8 : cy + 8, cx - 8 : cx - 8 + t] = 1
        img[cy - 8 : cy + 8, cx + 8 - t : cx + 8] = 1
        img[cy - 8 : cy - 8 + t, cx - 8 : cx + 8] = 1
        img[cy + 8 - t : cy + 8, cx - 8 : cx + 8] = 1
    elif cls == 8:  # solid box
        img[cy - 6 : cy + 6, cx - 6 : cx + 6] = 1
    else:  # double bar
        img[cy - 5 : cy - 5 + t, cx - 8 : cx + 8] = 1
        img[cy + 5 : cy + 5 + t, cx - 8 : cx + 8] = 1
    return img


def make_shape_dataset(n_samples=2000, seed=42):
    """Balanced dataset of the ten shapes, with noise, shuffled."""
    rng = np.random.default_rng(seed)
    X = np.empty((n_samples, 1, 28, 28), dtype=np.float32)
    y = np.empty(n_samples, dtype=np.int64)
    for i in range(n_samples):
        cls = i % 10
        X[i, 0] = np.clip(draw_shape(cls, rng) + rng.normal(0, 0.12, (28, 28)), 0, 1)
        y[i] = cls
    order = rng.permutation(n_samples)
    return torch.tensor(X[order]), torch.tensor(y[order])


class MNIST_CNN(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 16, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),  # 28x28 -> 14x14
            nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),  # 14x14 -> 7x7
        )
        # 32 * 7 * 7 is not a magic number: it is the shape the block above emits.
        # Get it wrong and you get "mat1 and mat2 shapes cannot be multiplied".
        self.classifier = nn.Sequential(
            nn.Linear(32 * 7 * 7, 64),
            nn.ReLU(),
            nn.Linear(64, num_classes),  # raw logits — CrossEntropyLoss wants logits
        )

    def forward(self, x):
        x = self.features(x)
        x = torch.flatten(x, 1)  # keep the batch dimension, flatten the rest
        return self.classifier(x)


def print_sample(img, label):
    """Show one training image as text, so you can see what the network sees."""
    print(f'\nsample input — class {label} ({CLASS_NAMES[label]}):')
    for row in range(4, 26):
        print('  ' + ''.join('#' if v > 0.5 else ('.' if v > 0.25 else ' ') for v in img[row]))


def train():
    torch.manual_seed(42)

    X, y = make_shape_dataset(2000)
    split = 1600
    X_train, y_train = X[:split], y[:split]
    X_val, y_val = X[split:], y[split:]
    print(f'train {tuple(X_train.shape)} | validation {tuple(X_val.shape)}')
    print_sample(X_train[0, 0].numpy(), int(y_train[0]))

    model = MNIST_CNN()
    n_params = sum(p.numel() for p in model.parameters())
    print(f'\nmodel has {n_params:,} parameters')

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    epochs, batch_size = 30, 64
    print('\n=== training ===')
    for epoch in range(1, epochs + 1):
        model.train()
        running = 0.0
        # Mini-batches, not one giant batch: more update steps per pass over the
        # data, which is how real training is always done.
        for start in range(0, len(X_train), batch_size):
            xb = X_train[start : start + batch_size]
            yb = y_train[start : start + batch_size]
            optimizer.zero_grad()
            loss = criterion(model(xb), yb)
            loss.backward()
            optimizer.step()
            running += loss.item() * len(xb)

        if epoch % 5 == 0 or epoch == 1:
            model.eval()
            with torch.no_grad():
                val_acc = (model(X_val).argmax(1) == y_val).float().mean().item() * 100
            print(f'epoch {epoch:2d}/{epochs} | train loss {running / len(X_train):.4f} | val accuracy {val_acc:.2f}%')

    model.eval()
    with torch.no_grad():
        preds = model(X_val).argmax(1)
    acc = (preds == y_val).float().mean().item() * 100
    print(f'\n=== final validation accuracy: {acc:.2f}%  (chance level is 10%) ===')

    # Per-class accuracy: an overall number can hide one class the model never gets.
    print('\nper-class accuracy:')
    for cls in range(10):
        mask = y_val == cls
        if mask.sum():
            hit = (preds[mask] == cls).float().mean().item() * 100
            print(f'  {cls} {CLASS_NAMES[cls]:16} {hit:6.1f}%  ({int(mask.sum())} samples)')


if __name__ == '__main__':
    train()

train (1600, 1, 28, 28) | validation (400, 1, 28, 28)

sample input — class 9 (double bar):
  ...
  #                        #
  ################
  ################
  ...

model has 105,866 parameters

=== training ===
epoch  1/30 | train loss 1.9055 | val accuracy 71.25%
epoch  5/30 | train loss 0.0070 | val accuracy 100.00%
epoch 10/30 | train loss 0.0011 | val accuracy 100.00%
...
epoch 30/30 | train loss 0.0001 | val accuracy 100.00%

=== final validation accuracy: 100.00%  (chance level is 10%) ===

per-class accuracy:
  0 ring              100.0%  (47 samples)
  1 vertical bar      100.0%  (31 samples)
  2 horizontal bar    100.0%  (35 samples)
  3 plus              100.0%  (38 samples)
  4 diagonal \        100.0%  (43 samples)
  5 diagonal /        100.0%  (41 samples)
  6 cross X           100.0%  (43 samples)
  7 hollow box        100.0%  (37 samples)
  8 solid box         100.0%  (47 samples)
  9 double bar        100.0%  (38 samples)