Residual Network (ResNet)

"""A residual block, forward pass only: y = relu(F(x) + x), written out.

Tensors are nested lists indexed [channel][row][col]. The convolution is the
literal sum from the CNN entry with a channel loop added and zero padding, so
the output keeps the input's spatial size. That matters: the shortcut can only
be added to F(x) if the two have exactly the same shape.
"""

import math


def conv3x3(x, weights):
    """x[c_in][h][w], weights[c_out][c_in][3][3]. Stride 1, zero padding 1."""
    c_in, h, w = len(x), len(x[0]), len(x[0][0])
    c_out = len(weights)
    out = [[[0.0] * w for _ in range(h)] for _ in range(c_out)]

    for oc in range(c_out):
        for i in range(h):
            for j in range(w):
                total = 0.0
                for ic in range(c_in):
                    for u in range(3):
                        for v in range(3):
                            row, col = i + u - 1, j + v - 1
                            if 0 <= row < h and 0 <= col < w:
                                total += x[ic][row][col] * weights[oc][ic][u][v]
                out[oc][i][j] = total
    return out


def conv1x1(x, weights):
    """The projection shortcut. weights[c_out][c_in]; mixes channels only."""
    c_in, h, w = len(x), len(x[0]), len(x[0][0])
    c_out = len(weights)
    return [
        [[sum(weights[oc][ic] * x[ic][i][j] for ic in range(c_in)) for j in range(w)]
         for i in range(h)]
        for oc in range(c_out)
    ]


def relu(x):
    return [[[max(0.0, value) for value in row] for row in plane] for plane in x]


def add(a, b):
    return [
        [[a[c][i][j] + b[c][i][j] for j in range(len(a[0][0]))] for i in range(len(a[0]))]
        for c in range(len(a))
    ]


def residual_block(x, w1, w2, w_projection=None):
    """y = relu(F(x) + shortcut(x)), with F = conv -> relu -> conv.

    The shortcut is x itself, unless the channel count changes, in which case a
    1x1 convolution projects x to the new shape so the addition is defined.
    """
    f = conv3x3(relu(conv3x3(x, w1)), w2)
    shortcut = x if w_projection is None else conv1x1(x, w_projection)
    return relu(add(f, shortcut))


def chain_gradient(weight, depth, is_residual):
    """d(output)/d(input) through a chain of scalar layers.

    A plain layer is y = weight * x, with derivative weight. A residual layer is
    y = weight * x + x, with derivative 1 + weight. The chain rule multiplies
    one factor per layer, so a small weight shrinks the plain product toward
    zero while the residual product has the identity term in every factor.
    """
    product = 1.0
    for _ in range(depth):
        product *= (1.0 + weight) if is_residual else weight
    return product


def global_average_pool(x):
    """Collapse each channel's map to its mean: the penultimate-layer embedding."""
    return [sum(sum(row) for row in plane) / (len(plane) * len(plane[0])) for plane in x]


def cosine(a, b):
    """Compare two embeddings by direction, ignoring their length."""
    dot = sum(p * q for p, q in zip(a, b))
    norm_a = math.sqrt(sum(p * p for p in a))
    norm_b = math.sqrt(sum(q * q for q in b))
    if norm_a == 0.0 or norm_b == 0.0:
        return 0.0
    return dot / (norm_a * norm_b)

"""ResNet - basic and bottleneck blocks, projection shortcuts, and the
initialisation details that make deep stacks train."""

from dataclasses import dataclass

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

BASIC_EXPANSION = 1
BOTTLENECK_EXPANSION = 4
STAGE_STRIDES_AFTER_FIRST = 2


@dataclass(frozen=True)
class ResNetConfig:
    block: str = "bottleneck"            # "basic" (18/34) or "bottleneck" (50+)
    layers: tuple[int, ...] = (3, 4, 6, 3)   # ResNet-50
    num_classes: int = 1000
    in_channels: int = 3
    base_width: int = 64
    zero_init_residual: bool = True

    def __post_init__(self) -> None:
        if self.block not in ("basic", "bottleneck"):
            raise ValueError(f"block must be 'basic' or 'bottleneck', got {self.block!r}")
        if not self.layers or any(count < 1 for count in self.layers):
            raise ValueError(f"every stage needs at least one block, got {self.layers}")


def _conv_bn(in_ch: int, out_ch: int, kernel: int, stride: int = 1) -> list[nn.Module]:
    """Conv -> BN with no conv bias: BN re-centres, so a bias would be wasted."""
    return [
        nn.Conv2d(in_ch, out_ch, kernel, stride=stride, padding=kernel // 2, bias=False),
        nn.BatchNorm2d(out_ch),
    ]


def _shortcut(in_ch: int, out_ch: int, stride: int) -> nn.Module:
    """Identity when shapes match; a 1x1 projection when channels or stride change."""
    if stride == 1 and in_ch == out_ch:
        return nn.Identity()
    return nn.Sequential(*_conv_bn(in_ch, out_ch, kernel=1, stride=stride))


class ResidualBlock(nn.Module):
    """y = relu(F(x) + shortcut(x)). The body F always ends in a BatchNorm."""

    def __init__(self, body: nn.Sequential, shortcut: nn.Module) -> None:
        super().__init__()
        self.body = body
        self.shortcut = shortcut
        self.act = nn.ReLU(inplace=True)

    @property
    def last_bn(self) -> nn.BatchNorm2d:
        last = self.body[-1]
        if not isinstance(last, nn.BatchNorm2d):
            raise TypeError("a residual body must end in BatchNorm2d")
        return last

    def forward(self, x: Tensor) -> Tensor:
        return self.act(self.body(x) + self.shortcut(x))


def basic_block(in_ch: int, width: int, stride: int) -> tuple[ResidualBlock, int]:
    """Two 3x3 convolutions. Used by ResNet-18 and ResNet-34."""
    out_ch = width * BASIC_EXPANSION
    body = nn.Sequential(
        *_conv_bn(in_ch, width, 3, stride), nn.ReLU(inplace=True), *_conv_bn(width, out_ch, 3)
    )
    return ResidualBlock(body, _shortcut(in_ch, out_ch, stride)), out_ch


def bottleneck_block(in_ch: int, width: int, stride: int) -> tuple[ResidualBlock, int]:
    """1x1 reduce -> 3x3 -> 1x1 expand. The 3x3 runs on few channels, which is
    where the roughly 17-fold parameter saving over a basic block comes from.
    Stride sits on the 3x3 (the v1.5 convention), not the first 1x1."""
    out_ch = width * BOTTLENECK_EXPANSION
    body = nn.Sequential(
        *_conv_bn(in_ch, width, 1),
        nn.ReLU(inplace=True),
        *_conv_bn(width, width, 3, stride),
        nn.ReLU(inplace=True),
        *_conv_bn(width, out_ch, 1),
    )
    return ResidualBlock(body, _shortcut(in_ch, out_ch, stride)), out_ch


class ResNet(nn.Module):
    def __init__(self, config: ResNetConfig) -> None:
        super().__init__()
        make_block = basic_block if config.block == "basic" else bottleneck_block

        self.stem = nn.Sequential(
            *_conv_bn(config.in_channels, config.base_width, 7, stride=2),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(3, stride=2, padding=1),
        )

        stages: list[nn.Module] = []
        in_ch = config.base_width
        for stage_index, block_count in enumerate(config.layers):
            width = config.base_width * 2**stage_index
            first_stride = 1 if stage_index == 0 else STAGE_STRIDES_AFTER_FIRST
            blocks: list[nn.Module] = []
            for block_index in range(block_count):
                block, in_ch = make_block(in_ch, width, first_stride if block_index == 0 else 1)
                blocks.append(block)
            stages.append(nn.Sequential(*blocks))
        self.stages = nn.Sequential(*stages)

        # Adaptive pooling makes the head independent of input resolution.
        self.pool = nn.AdaptiveAvgPool2d(1)
        self.head = nn.Linear(in_ch, config.num_classes)
        self._initialise(config.zero_init_residual)

    def _initialise(self, zero_init_residual: bool) -> None:
        for module in self.modules():
            if isinstance(module, nn.Conv2d):
                # He initialisation: variance 2 / fan keeps ReLU activations at a
                # stable scale from layer to layer.
                nn.init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu")
            elif isinstance(module, nn.BatchNorm2d):
                nn.init.ones_(module.weight)
                nn.init.zeros_(module.bias)
        if zero_init_residual:
            # Every block starts as exactly the identity, so the network begins
            # shallow and effective depth grows as the branches learn.
            for module in self.modules():
                if isinstance(module, ResidualBlock):
                    nn.init.zeros_(module.last_bn.weight)

    def features(self, x: Tensor) -> Tensor:
        """The penultimate-layer embedding: global-average-pooled final stage."""
        if x.dim() != 4:
            raise ValueError(f"expected (batch, channels, height, width), got {tuple(x.shape)}")
        return self.pool(self.stages(self.stem(x))).flatten(1)

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


@torch.no_grad()
def embed(model: ResNet, images: Tensor) -> Tensor:
    """L2-normalised embeddings, so a dot product between two rows is their cosine.

    eval() is not cosmetic: it switches BatchNorm to its running statistics.
    Embeddings from a model left in train mode depend on the batch they were
    computed in, which makes a vector index quietly inconsistent.
    """
    model.eval()
    return F.normalize(model.features(images), dim=1)

"""ResNet - channels-last, bf16, folded BatchNorm, compiled epilogues, batched retrieval."""

import torch
from torch import Tensor, nn


def fuse_conv_bn(conv: nn.Conv2d, bn: nn.BatchNorm2d) -> nn.Conv2d:
    """Fold an eval-mode BatchNorm into the convolution before it.

    BN at inference is a fixed per-channel affine map, so it merges into the
    weights: W' = W * gamma / sqrt(var + eps), b' = beta - mean * gamma / sqrt(var + eps).
    Numerically the same, one fewer layer and one fewer pass over the activations.
    """
    scale = bn.weight / torch.sqrt(bn.running_var + bn.eps)
    fused = nn.Conv2d(
        conv.in_channels, conv.out_channels, conv.kernel_size,
        stride=conv.stride, padding=conv.padding, bias=True,
    ).to(conv.weight.device)
    with torch.no_grad():
        fused.weight.copy_(conv.weight * scale.reshape(-1, 1, 1, 1))
        fused.bias.copy_(bn.bias - bn.running_mean * scale)
    return fused


def fuse_sequential(sequence: nn.Sequential) -> nn.Sequential:
    """Replace every Conv2d -> BatchNorm2d pair with one conv and an Identity."""
    layers = list(sequence.children())
    fused: list[nn.Module] = []
    index = 0
    while index < len(layers):
        has_partner = index + 1 < len(layers)
        if has_partner and isinstance(layers[index], nn.Conv2d) and isinstance(
            layers[index + 1], nn.BatchNorm2d
        ):
            fused.append(fuse_conv_bn(layers[index], layers[index + 1]))
            index += 2
        else:
            fused.append(layers[index])
            index += 1
    return nn.Sequential(*fused)


@torch.no_grad()
def prepare_for_inference(model: nn.Module, device: str) -> nn.Module:
    """Eval mode, folded BN in the stem, bodies and projection shortcuts, channels-last."""
    model = model.eval().to(device)
    model.stem = fuse_sequential(model.stem)
    # Collect first: mutating modules while iterating modules() is undefined.
    blocks = [module for module in model.modules() if hasattr(module, "body")]
    for block in blocks:
        block.body = fuse_sequential(block.body)
        if isinstance(block.shortcut, nn.Sequential):
            block.shortcut = fuse_sequential(block.shortcut)

    # Tensor cores want NHWC. Converting the weights once removes the
    # NCHW->NHWC transpose cuDNN would otherwise insert before each convolution.
    model = model.to(memory_format=torch.channels_last)
    # The compiler fuses the residual add and ReLU into the conv epilogue, so
    # the pre-activation sum is never written to memory.
    return torch.compile(model)


def train_step(model: nn.Module, optimizer, images: Tensor, labels: Tensor, device: str) -> float:
    """One step in bf16 autocast. bf16 keeps fp32's exponent range, so unlike
    fp16 it needs no loss scaling."""
    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)
    loss.backward()
    optimizer.step()
    return float(loss.detach())


@torch.inference_mode()
def build_index(model: nn.Module, loader, device: str) -> Tensor:
    """Embed a whole corpus in large batches: one L2-normalised matrix, kept in float32."""
    chunks = []
    for images in loader:
        images = images.to(device, non_blocking=True, memory_format=torch.channels_last)
        with torch.autocast(device_type=device, dtype=torch.bfloat16):
            features = model.features(images)
        # Normalise in float32: bf16 has an 8-bit mantissa, so near-duplicates
        # would tie and the ranking among them would be arbitrary.
        chunks.append(nn.functional.normalize(features.float(), dim=1))
    return torch.cat(chunks)


@torch.inference_mode()
def top_k(index: Tensor, queries: Tensor, k: int) -> tuple[Tensor, Tensor]:
    """Cosine nearest neighbours as one matrix product.

    Both sides are unit length, so queries @ index.T IS the cosine matrix: a
    single GEMM replaces a Python loop over query/item pairs.
    """
    scores = queries @ index.T
    return scores.topk(min(k, index.shape[0]), dim=1)