Ch 3 — Convolutional Neural Networks¶

# Load pretrained backbone
backbone = torchvision.models.resnet50(weights="IMAGENET1K_V2")
# Freeze backbone
for param in backbone.parameters():
    param.requires_grad = False
# Replace head
in_features = backbone.fc.in_features
backbone.fc = nn.Linear(in_features, num_classes)  # Only this trains

# Unfreeze last block only (layer4 in ResNet)
for param in backbone.layer4.parameters():
    param.requires_grad = True

# Different LR for head vs backbone
optimiser = torch.optim.AdamW([
    {"params": backbone.fc.parameters(), "lr": 1e-3},
    {"params": backbone.layer4.parameters(), "lr": 1e-4},
])

import torchvision.transforms as T

train_transform = T.Compose([
    T.RandomResizedCrop(224, scale=(0.08, 1.0)),  # Random crop
    T.RandomHorizontalFlip(p=0.5),                # Mirror
    T.ColorJitter(brightness=0.4, contrast=0.4,   # Colour
                  saturation=0.4, hue=0.1),
    T.ToTensor(),
    T.Normalize(mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225]),        # ImageNet stats
])

"""
ResNet-style CNN with basic residual blocks, batch normalisation,
global average pooling, and transfer-learning-compatible design.

Supports:
- Variable number of stages and blocks
- Bottleneck or basic block
- Fine-tuning vs feature extraction
"""

from __future__ import annotations

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


# ---------------------------------------------------------------------------
# Building blocks
# ---------------------------------------------------------------------------

class BasicBlock(nn.Module):
    """Two-layer residual block for ResNet-18/34.

    Args:
        in_channels:  Number of input channels.
        out_channels: Number of output channels.
        stride:       Stride for the first convolution (used for downsampling).
    """

    expansion: int = 1

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        stride: int = 1,
    ) -> None:
        super().__init__()

        self.conv1 = nn.Conv2d(
            in_channels, out_channels, kernel_size=3,
            stride=stride, padding=1, bias=False
        )
        self.bn1 = nn.BatchNorm2d(out_channels)

        self.conv2 = nn.Conv2d(
            out_channels, out_channels, kernel_size=3,
            stride=1, padding=1, bias=False
        )
        self.bn2 = nn.BatchNorm2d(out_channels)

        # Projection shortcut when dimensions change
        self.shortcut: nn.Module
        if stride != 1 or in_channels != out_channels:
            self.shortcut = nn.Sequential(
                nn.Conv2d(
                    in_channels, out_channels, kernel_size=1,
                    stride=stride, bias=False
                ),
                nn.BatchNorm2d(out_channels),
            )
        else:
            self.shortcut = nn.Identity()

    def forward(self, x: Tensor) -> Tensor:
        identity = self.shortcut(x)

        out = F.relu(self.bn1(self.conv1(x)), inplace=True)
        out = self.bn2(self.conv2(out))

        out = out + identity          # Skip connection: F(x) + x
        return F.relu(out, inplace=True)


class BottleneckBlock(nn.Module):
    """Three-layer bottleneck residual block for ResNet-50/101/152.

    Args:
        in_channels:  Number of input channels.
        out_channels: Number of output channels (expanded by factor 4).
        stride:       Stride for the 3×3 convolution.
    """

    expansion: int = 4

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        stride: int = 1,
    ) -> None:
        super().__init__()
        bottleneck_channels = out_channels // self.expansion

        self.conv1 = nn.Conv2d(in_channels, bottleneck_channels,
                               kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(bottleneck_channels)

        self.conv2 = nn.Conv2d(bottleneck_channels, bottleneck_channels,
                               kernel_size=3, stride=stride,
                               padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(bottleneck_channels)

        self.conv3 = nn.Conv2d(bottleneck_channels, out_channels,
                               kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(out_channels)

        self.shortcut: nn.Module
        if stride != 1 or in_channels != out_channels:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_channels, out_channels,
                          kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(out_channels),
            )
        else:
            self.shortcut = nn.Identity()

    def forward(self, x: Tensor) -> Tensor:
        identity = self.shortcut(x)

        out = F.relu(self.bn1(self.conv1(x)), inplace=True)
        out = F.relu(self.bn2(self.conv2(out)), inplace=True)
        out = self.bn3(self.conv3(out))

        out = out + identity
        return F.relu(out, inplace=True)


# ---------------------------------------------------------------------------
# Full ResNet
# ---------------------------------------------------------------------------

class ResNet(nn.Module):
    """Configurable ResNet.

    Args:
        block:       Block class (BasicBlock or BottleneckBlock).
        layers:      Number of blocks per stage, e.g. [2, 2, 2, 2] for ResNet-18.
        num_classes: Number of output classes.
        in_channels: Number of input image channels (default 3 for RGB).
    """

    def __init__(
        self,
        block: type[BasicBlock | BottleneckBlock],
        layers: list[int],
        num_classes: int = 1000,
        in_channels: int = 3,
    ) -> None:
        super().__init__()
        self._in_channels = 64

        # Stem
        self.stem = nn.Sequential(
            nn.Conv2d(in_channels, 64, kernel_size=7,
                      stride=2, padding=3, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1),
        )

        # Residual stages
        self.layer1 = self._make_stage(block, 64,  layers[0], stride=1)
        self.layer2 = self._make_stage(block, 128, layers[1], stride=2)
        self.layer3 = self._make_stage(block, 256, layers[2], stride=2)
        self.layer4 = self._make_stage(block, 512, layers[3], stride=2)

        # Head
        self.gap = nn.AdaptiveAvgPool2d((1, 1))          # Global average pool
        self.fc  = nn.Linear(512 * block.expansion, num_classes)

        self._init_weights()

    def _make_stage(
        self,
        block: type[BasicBlock | BottleneckBlock],
        out_channels: int,
        num_blocks: int,
        stride: int,
    ) -> nn.Sequential:
        """Build a stage of residual blocks."""
        blocks = [block(self._in_channels, out_channels, stride)]
        self._in_channels = out_channels
        for _ in range(1, num_blocks):
            blocks.append(block(out_channels, out_channels, stride=1))
        return nn.Sequential(*blocks)

    def _init_weights(self) -> None:
        """He initialisation for conv layers, zeros for BN biases."""
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode="fan_out",
                                        nonlinearity="relu")
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.ones_(m.weight)
                nn.init.zeros_(m.bias)

    def forward(self, x: Tensor) -> Tensor:
        """Forward pass returning logits.

        Args:
            x: Input images, shape (N, C, H, W).

        Returns:
            Logits, shape (N, num_classes).
        """
        x = self.stem(x)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = self.gap(x)
        x = torch.flatten(x, 1)
        return self.fc(x)


# ---------------------------------------------------------------------------
# Factory functions
# ---------------------------------------------------------------------------

def resnet18(num_classes: int = 1000) -> ResNet:
    """ResNet-18: BasicBlock, [2, 2, 2, 2]."""
    return ResNet(BasicBlock, [2, 2, 2, 2], num_classes=num_classes)


def resnet50(num_classes: int = 1000) -> ResNet:
    """ResNet-50: BottleneckBlock, [3, 4, 6, 3]."""
    return ResNet(BottleneckBlock, [3, 4, 6, 3], num_classes=num_classes)


# ---------------------------------------------------------------------------
# Demo
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    model = resnet18(num_classes=10)
    x = torch.randn(4, 3, 224, 224)
    logits = model(x)
    print(f"Output shape: {logits.shape}")   # (4, 10)

    total_params = sum(p.numel() for p in model.parameters())
    print(f"Parameters: {total_params:,}")  # ~11M