Example

The error names the classifier; the trace names the convolution

Check a derived output-shape rule against PyTorch over thousands of combinations, then remove one padding and find the first divergence.

A convolution changes two independent things: the channel count, which you declare, and the spatial extent, which you derive from the kernel, stride, padding and dilation. Get either wrong and the error arrives downstream, naming a layer that did nothing wrong.

This program composes the chapter’s pieces into one run: the derived rule checked against PyTorch over thousands of combinations, then a small network traced against a hand-written shape list, once as written and once with one padding removed.

Run it

Save this as cnn_geometry.py and run it with Python 3.11 or later. It needs PyTorch: pip install torch, then python cnn_geometry.py.

"""CNN geometry: derive the shape first, then find the first divergence."""
import math
from itertools import product

import torch
import torch.nn.functional as F
from torch import nn


def conv_out(size, k, s=1, p=0, d=1):
    return math.floor((size + 2 * p - d * (k - 1) - 1) / s + 1)


class SmallCNN(nn.Module):
    def __init__(self, conv2_padding=1):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 16, 3, padding=1)
        self.pool1 = nn.MaxPool2d(2)
        self.conv2 = nn.Conv2d(16, 32, 3, padding=conv2_padding)
        self.pool2 = nn.MaxPool2d(2)
        self.fc = nn.Linear(32 * 16 * 16, 10)

    def forward(self, x):
        x = self.pool1(F.relu(self.conv1(x)))
        x = self.pool2(F.relu(self.conv2(x)))
        return self.fc(torch.flatten(x, 1))


def ledger(model, x, expected):
    """Trace every leaf module's output shape and report the first divergence."""
    rows = []
    handles = [m.register_forward_hook(lambda mod, i, o: rows.append((type(mod).__name__, tuple(o.shape))))
               for m in model.modules() if not list(m.children())]
    error = None
    try:
        model(x)
    except RuntimeError as exc:
        error = str(exc).splitlines()[0]
    for handle in handles:
        handle.remove()
    first = None
    for (kind, shape), want in zip(rows, expected):
        if shape != want and first is None:
            first = kind
        print(f"  {kind:<10} expected={str(want):<16} observed={str(shape):<16} {'ok' if shape == want else 'DIVERGES'}")
    print("  first divergence:", first or "none")
    if error:
        print("  raised:", error)


tested = mismatches = 0
for size, k, s, p, d in product(range(4, 34), (1, 2, 3, 5, 7), (1, 2, 3), (0, 1, 2, 3), (1, 2, 3)):
    if size + 2 * p < d * (k - 1) + 1:
        continue
    derived = conv_out(size, k, s, p, d)
    observed = F.conv2d(torch.zeros(1, 1, size, size), torch.zeros(1, 1, k, k),
                        stride=s, padding=p, dilation=d).shape[-1]
    tested += 1
    mismatches += derived != observed
print(f"combinations tested: {tested}, mismatches: {mismatches}")

x = torch.randn(8, 3, 64, 64)
expected = [(8, 16, 64, 64), (8, 16, 32, 32), (8, 32, 32, 32), (8, 32, 16, 16), (8, 10)]
print("as written:")
ledger(SmallCNN(), x, expected)
print("conv2 padding removed:")
ledger(SmallCNN(conv2_padding=0), x, expected)

Expected outcome

combinations tested: 5055, mismatches: 0
as written:
  Conv2d     expected=(8, 16, 64, 64)  observed=(8, 16, 64, 64)  ok
  MaxPool2d  expected=(8, 16, 32, 32)  observed=(8, 16, 32, 32)  ok
  Conv2d     expected=(8, 32, 32, 32)  observed=(8, 32, 32, 32)  ok
  MaxPool2d  expected=(8, 32, 16, 16)  observed=(8, 32, 16, 16)  ok
  Linear     expected=(8, 10)          observed=(8, 10)          ok
  first divergence: none
conv2 padding removed:
  Conv2d     expected=(8, 16, 64, 64)  observed=(8, 16, 64, 64)  ok
  MaxPool2d  expected=(8, 16, 32, 32)  observed=(8, 16, 32, 32)  ok
  Conv2d     expected=(8, 32, 32, 32)  observed=(8, 32, 30, 30)  DIVERGES
  MaxPool2d  expected=(8, 32, 16, 16)  observed=(8, 32, 15, 15)  DIVERGES
  first divergence: Conv2d
  raised: mat1 and mat2 shapes cannot be multiplied (8x7200 and 8192x10)

Derive, then compare

The spatial rule is one line of arithmetic:

import math

def conv_out(size, k, s=1, p=0, d=1):
    return math.floor((size + 2*p - d*(k - 1) - 1) / s + 1)

The sweep checks the rule against PyTorch: 5,055 combinations, no mismatches.

The ledger is the other half of the technique: write the expected shape after every layer, run the model, and compare. Remove one padding from conv2 and the exception names the Linear layer, where the mismatch became illegal. The ledger names conv2, where the geometry changed.

Mechanism and limitations

Documented: the chapter derives the output rule, sweeps input sizes 4 to 33 against kernels, strides, paddings and dilations for 5055 combinations with 0 mismatches, and lists a network’s expected shapes.

Observed: the run above. The sweep tested 5,055 combinations, no mismatches. The network as written matched every expected shape, with no divergence. With conv2’s padding removed, conv2 produced (8, 32, 30, 30) where (8, 32, 32, 32) was expected, and the run ended in a matrix error.

Not shown. The program composes the chapter’s fragments into one run: it adds the ledger printer and sweeps with F.conv2d on zero tensors. It checks one network at one resolution, and does not show that a ledger proves dimensions compose rather than that the axes mean what you meant.

Canonical source , the chapter this example composes. Read the chapter .

Understand this example

Copy this prompt into your AI tool. No code runs here.

Apply this example

Copy this prompt into your AI tool. No code runs here.