Convolutional Neural Network

"""2-D convolution and a forward pass - the operation, written out.

Four nested loops: output position (i, j), then kernel offset (u, v). This is
literally the sum in the equation above.
"""


def conv2d(image, kernel, stride=1):
    """Single-channel valid convolution. image[h][w], kernel[kh][kw]."""
    h, w = len(image), len(image[0])
    kh, kw = len(kernel), len(kernel[0])
    out_h = (h - kh) // stride + 1
    out_w = (w - kw) // stride + 1

    output = [[0.0] * out_w for _ in range(out_h)]

    for i in range(out_h):
        for j in range(out_w):
            total = 0.0
            # The SAME kernel weights are used at every (i, j) - that reuse is
            # weight sharing, and it is the entire point of the architecture.
            for u in range(kh):
                for v in range(kw):
                    total += image[i * stride + u][j * stride + v] * kernel[u][v]
            output[i][j] = total

    return output


def relu(feature_map):
    return [[max(0.0, v) for v in row] for row in feature_map]


def max_pool(feature_map, size=2):
    """Downsample by taking the strongest activation in each window."""
    h, w = len(feature_map), len(feature_map[0])
    out = [[0.0] * (w // size) for _ in range(h // size)]

    for i in range(h // size):
        for j in range(w // size):
            window = [
                feature_map[i * size + u][j * size + v]
                for u in range(size)
                for v in range(size)
            ]
            out[i][j] = max(window)

    return out


def forward(image, kernels):
    """Conv -> ReLU -> pool, once per kernel. Receptive field widens with depth."""
    return [max_pool(relu(conv2d(image, k))) for k in kernels]

"""CNN - typed, validated, with the train/eval distinction made explicit."""

from dataclasses import dataclass

import torch
from torch import Tensor, nn


@dataclass(frozen=True)
class CnnConfig:
    in_channels: int = 3
    num_classes: int = 10
    base_width: int = 32
    depth: int = 3
    dropout: float = 0.1


class ConvBlock(nn.Module):
    """Conv -> BatchNorm -> ReLU, the standard unit.

    Bias is disabled on the convolution because BatchNorm immediately
    re-centres the activations - the bias term would be learned and then
    subtracted away, wasting parameters and slightly hurting conditioning.
    """

    def __init__(self, in_ch: int, out_ch: int) -> None:
        super().__init__()
        self.block = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
        )

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


class SimpleCnn(nn.Module):
    def __init__(self, config: CnnConfig) -> None:
        super().__init__()
        if config.depth < 1:
            raise ValueError(f"depth must be >= 1, got {config.depth}")
        if not 0.0 <= config.dropout < 1.0:
            raise ValueError(f"dropout must be in [0, 1), got {config.dropout}")

        layers: list[nn.Module] = []
        in_ch = config.in_channels
        width = config.base_width

        for _ in range(config.depth):
            layers.append(ConvBlock(in_ch, width))
            layers.append(nn.MaxPool2d(2))
            in_ch = width
            width *= 2      # channels double as resolution halves

        self.features = nn.Sequential(*layers)
        # Adaptive pooling makes the head independent of input resolution -
        # a fixed Linear would break the moment the image size changed.
        self.pool = nn.AdaptiveAvgPool2d(1)
        self.head = nn.Sequential(
            nn.Dropout(config.dropout),
            nn.Linear(in_ch, config.num_classes),
        )

    def forward(self, x: Tensor) -> Tensor:
        if x.dim() != 4:
            raise ValueError(f"expected (batch, channels, height, width), got {tuple(x.shape)}")
        features = self.pool(self.features(x)).flatten(1)
        return self.head(features)


@torch.no_grad()
def evaluate(model: nn.Module, loader, device: str) -> float:
    """Accuracy on a loader.

    model.eval() is not cosmetic: it switches BatchNorm from batch statistics
    to running averages and disables Dropout. Evaluating in train mode gives
    wrong and batch-size-dependent results.
    """
    model.eval()
    correct = total = 0

    for images, labels in loader:
        images, labels = images.to(device), labels.to(device)
        predictions = model(images).argmax(dim=1)
        correct += int((predictions == labels).sum())
        total += labels.numel()

    return correct / max(total, 1)

"""CNN - channels-last, mixed precision, fused conv-bn for inference."""

import torch
from torch import Tensor, nn


def prepare_for_training(model: nn.Module, device: str) -> nn.Module:
    """channels_last + bfloat16 autocast.

    Tensor cores want NHWC. PyTorch defaults to NCHW, so cuDNN inserts a
    transpose before every convolution; converting once up front removes that
    transpose from every layer of every step.
    """
    model = model.to(device, memory_format=torch.channels_last)
    return model


def train_epoch(model: nn.Module, optimizer, loader, device: str) -> float:
    model.train()
    scaler = torch.amp.GradScaler(device)
    total_loss = 0.0

    for images, labels in loader:
        # non_blocking with pinned memory overlaps the host-to-device copy
        # with compute, so the GPU is not idling between batches.
        images = images.to(device, non_blocking=True, memory_format=torch.channels_last)
        labels = labels.to(device, non_blocking=True)

        optimizer.zero_grad(set_to_none=True)
        with torch.autocast(device_type=device, dtype=torch.bfloat16):
            loss = nn.functional.cross_entropy(model(images), labels)

        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()
        total_loss += float(loss.item())

    return total_loss / max(len(loader), 1)


@torch.no_grad()
def fuse_conv_bn(conv: nn.Conv2d, bn: nn.BatchNorm2d) -> nn.Conv2d:
    """Fold BatchNorm into the preceding convolution for inference.

    At eval time BN is a fixed affine transform, so it can be folded into the
    conv weights algebraically:
        W' = W * gamma / sqrt(var + eps)
        b' = (b - mean) * gamma / sqrt(var + eps) + beta
    The result is numerically identical and removes an entire layer - both its
    kernel launch and its full pass over the activations - from every forward.
    """
    fused = nn.Conv2d(
        conv.in_channels, conv.out_channels, conv.kernel_size,
        conv.stride, conv.padding, bias=True,
    )

    scale = bn.weight / torch.sqrt(bn.running_var + bn.eps)
    fused.weight.copy_(conv.weight * scale.reshape(-1, 1, 1, 1))

    conv_bias = conv.bias if conv.bias is not None else torch.zeros(conv.out_channels)
    fused.bias.copy_((conv_bias - bn.running_mean) * scale + bn.bias)

    return fused