Image Classification with Transfer Learning in PyTorch: From a Small CNN to a Fine-Tuned ResNet

"""Harrow County maps land use from satellite tiles. A surveyor hand-labelled 1,000 of them; this file turns
those labels into a classifier for the rest. Tiles are 64x64 RGB Sentinel-2 images (EuroSAT, Helber et al.)."""
import numpy as np
import torch
import torchvision
from sklearn.model_selection import train_test_split
from torch import nn

torch.set_num_threads(1)   # the lab has one CPU; more threads only fight over it

DATA = "/opt/datasets/eurosat-2500.npz"
WEIGHTS = "/opt/weights/resnet18-f37072fd.pth"     # ResNet18 trained on ImageNet's 1.28 million photos
MEAN = torch.tensor([0.485, 0.456, 0.406])         # ImageNet's per-channel pixel mean and std (0-1 scale):
STD = torch.tensor([0.229, 0.224, 0.225])          # a pretrained network expects its inputs scaled this way
CUT = "layer4"


def load_data():
    """{"train": (images, labels), "val": ..., "test": ...}: 1,000 / 500 / 1,000 tiles, stratified. Images are
    uint8 numpy arrays of shape (n, 64, 64, 3); labels int64. Also returns the class names."""
    d = np.load(DATA)
    X, y = d["images"], d["labels"]
    idx = np.arange(len(y))
    tr, rest = train_test_split(idx, train_size=1000, stratify=y, random_state=0)
    va, te = train_test_split(rest, train_size=500, stratify=y[rest], random_state=0)
    return {"train": (X[tr], y[tr]), "val": (X[va], y[va]), "test": (X[te], y[te])}, [str(c) for c in d["classes"]]


# ---------- Step 1: tiles as tensors ----------
def to_tensor(images):
    """uint8 (n, 64, 64, 3) -> float32 (n, 3, 64, 64): channels first, scaled to 0-1, then normalised per
    channel with MEAN and STD."""
    # TODO (Step 1): permute to (n, 3, 64, 64), float, / 255, then subtract MEAN and divide by STD
    # (reshape both to (1, 3, 1, 1) so they apply per channel).
    raise NotImplementedError("Step 1: write to_tensor()")


def augment(x, generator):
    """A randomly turned copy of the batch: every tile independently gets a horizontal flip with probability 0.5
    and then a rotation by k * 90 degrees, k drawn from 0-3. Draw the flips first, then the k's, both from
    `generator`: flips = torch.rand(n, generator=generator) < 0.5, ks = torch.randint(0, 4, (n,), generator=generator)."""
    # TODO (Step 1): draw the flips, then the ks, from the generator; flip each tile (dims 2), then
    # torch.rot90(tile, k, dims=(1, 2)). Return a new tensor.
    raise NotImplementedError("Step 1: write augment()")


# ---------- Step 2: a small CNN from scratch ----------
class SmallCNN(nn.Module):
    """Three blocks of Conv2d(3x3, padding=1) -> BatchNorm2d -> ReLU -> MaxPool2d(2), with 32, 64 and 128
    channels (64x64 -> 32 -> 16 -> 8), then the mean over the 8x8 positions, Dropout(0.3) and Linear(128, 10)."""

    def __init__(self):
        super().__init__()
        raise NotImplementedError("SmallCNN.__init__() arrives in Step 2")

    def forward(self, x):
        raise NotImplementedError("SmallCNN.forward() arrives in Step 2")


def accuracy(model, X, y, batch_size=250):
    """Share of tiles predicted right, in evaluation mode, without gradients, in batches."""
    raise NotImplementedError("accuracy() arrives in Step 2")


def train_classifier(model, X, y, optimizer, epochs, batch_size=32, seed=0, turn=True, freeze_bn=False):
    """Mini-batch training with cross-entropy; each batch is augment()ed when turn=True. Returns epoch losses."""
    g = torch.Generator().manual_seed(seed)
    y = torch.as_tensor(y)
    losses = []
    for _ in range(epochs):
        set_mode(model, freeze_bn) if freeze_bn else model.train()
        perm = torch.randperm(len(X), generator=g)
        total = 0.0
        for k in range(0, len(X), batch_size):
            idx = perm[k:k + batch_size]
            xb = augment(X[idx], g) if turn else X[idx]
            optimizer.zero_grad()
            loss = nn.functional.cross_entropy(model(xb), y[idx])
            loss.backward()
            optimizer.step()
            total += loss.item() * len(idx)
        losses.append(total / len(X))
    return losses


# ---------- Step 3: a pretrained backbone, frozen ----------
LAYERS = ["conv1", "bn1", "relu", "maxpool", "layer1", "layer2", "layer3", "layer4"]


def backbone(cut):
    """The ImageNet-trained ResNet18 (weights from WEIGHTS) up to and including the block named `cut`, as an
    nn.Sequential of its children in LAYERS order; in eval mode, every parameter frozen (requires_grad False)."""
    raise NotImplementedError("backbone() arrives in Step 3")


def features(body, X, batch_size=100):
    """One vector per tile: body's output averaged over its spatial positions, (n, C, h, w) -> (n, C). Eval mode,
    no gradients, in batches."""
    raise NotImplementedError("features() arrives in Step 3")


class Standardize(nn.Module):
    """Fixed scaling learned from training features: (x - mean) / std."""

    def __init__(self, mean, std):
        super().__init__()
        self.register_buffer("mean", mean.clone())
        self.register_buffer("std", std.clone())

    def forward(self, x):
        return (x - self.mean) / self.std


def fit_probe(F, y, l2=1e-3):
    """A linear classifier on frozen features: nn.Sequential(Standardize(F.mean(0), F.std(0) + 1e-6),
    nn.Linear(C, 10)). After torch.manual_seed(0), fit the Linear with torch.optim.LBFGS(max_iter=200,
    line_search_fn="strong_wolfe") on the whole set at once, minimising cross-entropy + l2 * (weight ** 2).sum()."""
    raise NotImplementedError("fit_probe() arrives in Step 3")


# ---------- Step 4: where to cut ----------
def compare_cuts(cuts, data):
    """{cut: validation accuracy of fit_probe() trained on the training tiles' features from backbone(cut)}"""
    raise NotImplementedError("compare_cuts() arrives in Step 4")


# ---------- Step 5: fine-tune, carefully ----------
class TileClassifier(nn.Module):
    """Backbone -> mean over positions -> head. One model that takes tiles and returns 10 scores."""

    def __init__(self, body, head):
        super().__init__()
        self.body, self.head = body, head

    def forward(self, x):
        return self.head(self.body(x).mean(dim=(2, 3)))


def set_mode(model, freeze_bn=True):
    """model.train(), then every BatchNorm2d back to eval(): its running statistics came from 1.28 million
    photos and 32 tiles at a time would overwrite them."""
    raise NotImplementedError("set_mode() arrives in Step 5")


def prepare_finetune(model, lr_body=5e-5, lr_head=3e-4):
    """Unfreeze the last block of model.body (model.body[-1]) and the head's Linear; everything else stays
    frozen. Returns an Adam optimizer with two parameter groups: the last block at lr_body, the head at lr_head."""
    raise NotImplementedError("prepare_finetune() arrives in Step 5")