Residual Networks and Transfer Learning

"""DL Week 5: Residual Networks and Transfer Learning.

An independent teaching example accompanying the Algorithms II Week 6 material.
Run with: python -m pip install torch
          python dl-week-5.py

Pretrains a tiny residual backbone on a clean synthetic source task, then
compares three strategies on a smaller, noisier target task: a frozen backbone
with a new head, full fine-tuning, and training from scratch. Everything is
generated and trained locally; no checkpoint is downloaded.
"""
import argparse
import copy

import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

torch.set_num_threads(1)

SIZE = 16


def make_stripes(count, seed, noise, period=4):
    """Half horizontal, half vertical stripes at a random phase, plus noise."""
    g = torch.Generator().manual_seed(seed)
    images = torch.zeros(count, 1, SIZE, SIZE)
    labels = torch.zeros(count, dtype=torch.long)
    for i in range(count):
        vertical = i % 2
        phase = int(torch.randint(0, period, (1,), generator=g).item())
        band = ((torch.arange(SIZE) + phase) % period < period // 2).float()
        images[i, 0] = band.repeat(SIZE, 1) if vertical else band.unsqueeze(1).repeat(1, SIZE)
        labels[i] = vertical
    images += noise * torch.randn(images.shape, generator=g)
    return TensorDataset(images, labels)


class ResidualBlock(nn.Module):
    """y = relu(F(x) + shortcut(x)), with a 1x1 projection when shapes differ."""

    def __init__(self, c_in, c_out, stride=1):
        super().__init__()
        self.branch = nn.Sequential(
            nn.Conv2d(c_in, c_out, 3, stride=stride, padding=1, bias=False),
            nn.BatchNorm2d(c_out), nn.ReLU(),
            nn.Conv2d(c_out, c_out, 3, padding=1, bias=False),
            nn.BatchNorm2d(c_out),
        )
        same = stride == 1 and c_in == c_out
        self.shortcut = nn.Identity() if same else nn.Sequential(
            nn.Conv2d(c_in, c_out, 1, stride=stride, bias=False), nn.BatchNorm2d(c_out)
        )

    def forward(self, x):
        return nn.functional.relu(self.branch(x) + self.shortcut(x))


def make_backbone():
    return nn.Sequential(
        nn.Conv2d(1, 8, 3, padding=1, bias=False), nn.BatchNorm2d(8), nn.ReLU(),
        ResidualBlock(8, 8),
        ResidualBlock(8, 16, stride=2),
        nn.AdaptiveAvgPool2d(1), nn.Flatten(),
    )


def make_head():
    return nn.Linear(16, 2)


def evaluate(backbone, head, loader):
    backbone.eval(); head.eval()
    loss_sum = correct = total = 0
    with torch.no_grad():
        for images, labels in loader:
            logits = head(backbone(images))
            loss_sum += nn.functional.cross_entropy(logits, labels, reduction="sum").item()
            correct += (logits.argmax(1) == labels).sum().item()
            total += labels.numel()
    return loss_sum / total, correct / total


def train(backbone, head, train_loader, val_loader, epochs, groups, freeze_backbone):
    """Fit, restoring whichever parameters gave the lowest validation loss."""
    optimizer = torch.optim.AdamW(groups)
    best = float("inf")
    best_state = (copy.deepcopy(backbone.state_dict()), copy.deepcopy(head.state_dict()))
    for _ in range(epochs):
        # model.train() would put a frozen backbone's BatchNorm back into
        # training mode and let its running statistics drift, so set the modes
        # explicitly every epoch rather than once at the start.
        head.train()
        backbone.eval() if freeze_backbone else backbone.train()
        for images, labels in train_loader:
            optimizer.zero_grad()
            if freeze_backbone:
                with torch.no_grad():
                    features = backbone(images)
            else:
                features = backbone(images)
            loss = nn.functional.cross_entropy(head(features), labels)
            loss.backward()
            optimizer.step()
        val_loss, _ = evaluate(backbone, head, val_loader)
        if val_loss < best:
            best = val_loss
            best_state = (copy.deepcopy(backbone.state_dict()), copy.deepcopy(head.state_dict()))
    backbone.load_state_dict(best_state[0]); head.load_state_dict(best_state[1])
    return best


def buffer_signature(module):
    """A single number summarising every BatchNorm running statistic."""
    return sum(float(b.sum()) for b in module.buffers())


def counts(backbone, head):
    total = sum(p.numel() for p in list(backbone.parameters()) + list(head.parameters()))
    trainable = sum(p.numel() for p in list(backbone.parameters()) + list(head.parameters()) if p.requires_grad)
    return total, trainable


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--smoke", action="store_true", help="execution check, not a convergence result")
    args = parser.parse_args()

    torch.manual_seed(0)
    source_epochs = 2 if args.smoke else 14
    target_epochs = 2 if args.smoke else 24
    source_n = 128 if args.smoke else 1024
    target_n = 16 if args.smoke else 16

    # Source: many examples, and imagery that resembles the target. Target: very
    # few labels. That combination is the case where reusing features should pay.
    # Lower SOURCE_NOISE towards 0.25 to widen the domain gap and watch the
    # frozen backbone lose its advantage.
    SOURCE_NOISE, TARGET_NOISE = 0.9, 1.6
    source_train = DataLoader(make_stripes(source_n, 1, noise=SOURCE_NOISE), batch_size=32, shuffle=True)
    source_val = DataLoader(make_stripes(256, 2, noise=SOURCE_NOISE), batch_size=128)
    target_train = DataLoader(make_stripes(target_n, 3, noise=TARGET_NOISE), batch_size=8, shuffle=True)
    target_val = DataLoader(make_stripes(256, 4, noise=TARGET_NOISE), batch_size=128)
    target_test = DataLoader(make_stripes(512, 5, noise=TARGET_NOISE), batch_size=128)

    print("Shape check: the residual branch and its projection must agree")
    block = ResidualBlock(8, 16, stride=2)
    probe = torch.zeros(2, 8, 16, 16)
    with torch.no_grad():
        branch_out, short_out = block.branch(probe), block.shortcut(probe)
    print(f"  input     {tuple(probe.shape)}")
    print(f"  branch    {tuple(branch_out.shape)}")
    print(f"  shortcut  {tuple(short_out.shape)}")
    print(f"  addable   {branch_out.shape == short_out.shape}")

    print(f"\nPretraining the backbone on the clean source task ({source_n} examples)")
    backbone, head = make_backbone(), make_head()
    train(backbone, head, source_train, source_val, source_epochs,
          [{"params": list(backbone.parameters()) + list(head.parameters()), "lr": 3e-3}], False)
    _, source_acc = evaluate(backbone, head, source_val)
    print(f"  source validation accuracy {source_acc:.3f}")
    pretrained = copy.deepcopy(backbone.state_dict())

    print(f"\nTarget task: {target_n} noisy examples, three strategies")
    results = {}

    # 1. Frozen backbone, new head.
    b, h = make_backbone(), make_head()
    b.load_state_dict(pretrained)
    for p in b.parameters():
        p.requires_grad_(False)
    before = buffer_signature(b)
    total, trainable = counts(b, h)
    train(b, h, target_train, target_val, target_epochs, [{"params": h.parameters(), "lr": 5e-3}], True)
    after = buffer_signature(b)
    _, acc = evaluate(b, h, target_test)
    results["Frozen backbone"] = (total, trainable, acc)
    print(f"  frozen    buffers before {before:.6f} after {after:.6f} -> "
          f"{'unchanged' if abs(before - after) < 1e-9 else 'CHANGED'}")

    # 2. Fine-tune everything, with a smaller rate on the pretrained weights.
    b, h = make_backbone(), make_head()
    b.load_state_dict(pretrained)
    total, trainable = counts(b, h)
    train(b, h, target_train, target_val, target_epochs,
          [{"params": b.parameters(), "lr": 3e-4}, {"params": h.parameters(), "lr": 3e-3}], False)
    _, acc = evaluate(b, h, target_test)
    results["Fine-tuned"] = (total, trainable, acc)

    # 3. No source data at all.
    b, h = make_backbone(), make_head()
    total, trainable = counts(b, h)
    train(b, h, target_train, target_val, target_epochs,
          [{"params": list(b.parameters()) + list(h.parameters()), "lr": 3e-3}], False)
    _, acc = evaluate(b, h, target_test)
    results["From scratch"] = (total, trainable, acc)

    print(f"\n  {'strategy':<18}{'total':>9}{'trainable':>11}{'test acc':>10}")
    for name, (total, trainable, acc) in results.items():
        print(f"  {name:<18}{total:>9,}{trainable:>11,}{acc:>10.3f}")

    print("\nThe two transfer rows have already seen 1024 source examples; the scratch")
    print("row has only the 16 target labels. That is a difference in data budget, not")
    print("only in method, and it is the main reason the gap is large.")
    print("Try SOURCE_NOISE = 0.25 to make the source look unlike the target: the")
    print("frozen backbone can then fall to chance while fine-tuning still adapts.")
    print("One run on one synthetic pair — vary the seed before ranking anything.")
    if args.smoke:
        print("Smoke run: execution only. These numbers are not a convergence result.")


if __name__ == "__main__":
    main()

A correction F(x) to the shortcut representation, with the sum representing the desired transformation.

With a shape-preserving identity shortcut and zero residual branch, the sum is x; a post-sum activation can still modify negative coordinates.

It mixes channels and can subsample with stride to make the shortcut's dimensions match the residual branch.

eval() changes mode-dependent layers. no_grad() disables recording operations for gradient computation. They serve different purposes.

It begins without task-specific weights, while pretrained features may need smaller adjustments. Validate the choice.

Reusing source training hurts the target outcome relative to an appropriate baseline under the stated evaluation protocol.

It isolates the shortcut while controlling other architectural and training differences, then examines both training and held-out behavior.

When it changes the label, removes essential evidence, or creates unrealistic examples that undermine the target task.