Transfer Learning — Programuoki

Code

count_trainable and RGB pictures — write them yourself

build_resnet works. Two more things you cannot do without it — this time empty skeletons, no holes:

count_trainable — how many parameters are really trainable.
The RGB picture path in data.py — colour, 224 × 224, normalised the way ImageNet was.

1. `count_trainable`

It is one line, but it is a check at the start of every training run. By printing "trainable parameters: 2,052 of 11,178,564" you see whether freezing really worked, before you spend five minutes.

trainable ← number of parameters of those p whose p.requires_grad is True
total     ← number of all parameters
return (trainable, total)

The number of elements of a parameter p is p.numel().

2. RGB pictures

We added the parameters size (the side) and rgb (whether the picture is colour). Two things await you:

build_transform, the rgb=True branch. The picture has to be resized to (size, size) — transforms.Resize. Do not turn it gray again: ImageFolder already reads the file as 3-channel RGB (a gray PNG gets its one channel tripled).
make_loaders, the rgb=True branch. The mean and std come not from your data but from ImageNet (IMAGENET_MEAN, IMAGENET_STD, already at the top). A pretrained network expects that scale.

Note that the rgb=False (default) branch stays as it was — lesson 9's code works unchanged.

The test

test_transfer.py has grown. The new tests check that:

count_trainable returns (2052, 11178564) for a frozen network and (8395780, 11178564) with layer4;
the default loaders still give [N, 1, 28, 28];
rgb=True gives [N, 3, 64, 64];
the pixels agree with a hand-computed (x − mean) / std using the ImageNet statistics.

Run python -m unittest test_transfer (you already have the holes filled):

EE.....EE
======================================================================
ERROR: test_count_trainable_frozen_backbone (...)
NotImplementedError: return (trainable parameters, all parameters)

ERROR: test_count_trainable_with_layer4 (...)
NotImplementedError: return (trainable parameters, all parameters)

ERROR: test_rgb_loaders_give_three_channels_at_the_requested_size (...)
NotImplementedError: use IMAGENET_MEAN and IMAGENET_STD

ERROR: test_rgb_pixels_use_the_imagenet_statistics (...)
NotImplementedError: use IMAGENET_MEAN and IMAGENET_STD
----------------------------------------------------------------------
Ran 9 tests in 0.7s

FAILED (errors=4)

(Shortened.) Five tests pass, four wait. Once you fill in one make_loaders branch, you will see the second message: write the rgb branch.

transfer.py and data.py are shown with skeletons in the panel. The panel shows the skeleton, not your file.

When it does not pass

RuntimeError: output with shape [1, 64, 64] doesn't match the broadcast shape [3, 64, 64] — a Grayscale was left in the rgb branch: the picture has one channel, the ImageNet means have three.
Tensor-likes are not close! Greatest absolute difference: 0.92… — the pixels were normalised with other numbers (e.g. [0.5, 0.5, 0.5]).
Tuples differ: (11178564, 2052) != (2052, 11178564) — count_trainable returns (total, trainable) and it should be the other way round.

transfer.py MODIFY 4 lines changed

Edit this file in ./

"""transfer.py — stand on a pretrained ResNet-18 and teach it YOUR categories."""

from torch import nn
from torchvision.models import ResNet18_Weights, resnet18


def build_resnet(num_classes: int, finetune_last_block: bool = False) -> nn.Module:
    model = resnet18(weights=ResNet18_Weights.DEFAULT)  # downloads once (~45 MB)

    for p in model.parameters():  # freeze everything it already knows
        # HOLE 1: switch off gradient tracking for p (one line).
        pass
    if finetune_last_block:  # optionally let the last residual block adapt
        for p in model.layer4.parameters():
            p.requires_grad = True

    # HOLE 2: replace model.fc with a brand-new nn.Linear that maps the
    # network's feature vector (model.fc.in_features numbers) to num_classes.
    return model


def count_trainable(model: nn.Module) -> tuple[int, int]:
    raise NotImplementedError("return (trainable parameters, all parameters)")

data.py MODIFY 30 lines changed

Edit this file in ./

"""data.py — your own images, read from folders: data/<class>/*.png."""

from collections import Counter

import torch
from torch.utils.data import DataLoader, Subset, WeightedRandomSampler
from torchvision import datasets, transforms

# What the pretrained networks were trained with (ImageNet pixel statistics).
IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD = [0.229, 0.224, 0.225]


def class_counts(root: str) -> dict[str, int]:
    """How many images each class folder holds."""
    ds = datasets.ImageFolder(root)
    counts = Counter(ds.targets)
    return {name: counts[i] for i, name in enumerate(ds.classes)}


def build_transform(
    augment: bool,
    mean: list[float] | None = None,
    std: list[float] | None = None,
    size: int = 28,
    rgb: bool = False,
) -> transforms.Compose:
    if rgb:  # pretrained nets want 3 channels at a larger size
        # Start steps with a Resize to (size, size) and keep all 3 colour
        # channels - no Grayscale here.
        raise NotImplementedError("write the rgb branch")
    else:
        steps = [transforms.Grayscale(num_output_channels=1)]
    if augment:  # random changes, applied freshly every time an image is loaded
        steps += [
            transforms.RandomHorizontalFlip(),
            transforms.RandomAffine(degrees=10, translate=(0.1, 0.1), scale=(0.9, 1.1)),
        ]
    steps.append(transforms.ToTensor())  # PIL image -> [1, H, W] floats in 0..1
    if mean is not None:
        steps.append(transforms.Normalize(mean, std))
    return transforms.Compose(steps)


def split_indices(
    targets: list[int], val_frac: float, seed: int
) -> tuple[list[int], list[int]]:
    """A stratified split: every class contributes val_frac of its images."""
    gen = torch.Generator().manual_seed(seed)
    train_idx: list[int] = []
    val_idx: list[int] = []
    for cls in sorted(set(targets)):
        idx = [i for i, t in enumerate(targets) if t == cls]
        order = torch.randperm(len(idx), generator=gen).tolist()
        n_val = max(1, round(len(idx) * val_frac))
        val_idx += [idx[j] for j in order[:n_val]]
        train_idx += [idx[j] for j in order[n_val:]]
    return train_idx, val_idx


def mean_std(ds: datasets.ImageFolder, idx: list[int]) -> tuple[float, float]:
    """Pixel mean and std over the TRAINING images only."""
    pixels = torch.cat([ds[i][0].flatten() for i in idx])
    return pixels.mean().item(), pixels.std().item()


def make_loaders(
    root: str,
    batch_size: int = 32,
    val_frac: float = 0.2,
    augment: bool = False,
    balance: bool = False,
    seed: int = 42,
    size: int = 28,
    rgb: bool = False,
) -> tuple[DataLoader, DataLoader, list[str]]:
    """Return (train_loader, val_loader, class_names)."""
    plain = datasets.ImageFolder(root, transform=build_transform(augment=False))
    train_idx, val_idx = split_indices(plain.targets, val_frac, seed)
    if rgb:
        # Normalise with the statistics the pretrained network was trained with.
        raise NotImplementedError("use IMAGENET_MEAN and IMAGENET_STD")
    else:
        m, s = mean_std(plain, train_idx)
        mean, std = [m], [s]

    # Two views of the same files: augmented for training, untouched for validation.
    train_ds = datasets.ImageFolder(
        root, transform=build_transform(augment, mean, std, size, rgb)
    )
    val_ds = datasets.ImageFolder(
        root, transform=build_transform(False, mean, std, size, rgb)
    )

    sampler = None
    if balance:  # draw rare classes more often
        counts = Counter(plain.targets[i] for i in train_idx)
        weights = [1.0 / counts[plain.targets[i]] for i in train_idx]
        sampler = WeightedRandomSampler(weights, num_samples=len(train_idx))

    train_loader = DataLoader(
        Subset(train_ds, train_idx),
        batch_size=batch_size,
        shuffle=(sampler is None),
        sampler=sampler,
    )
    val_loader = DataLoader(Subset(val_ds, val_idx), batch_size=batch_size)
    return train_loader, val_loader, plain.classes

test_transfer.py MODIFY 49 lines changed

Edit this file in ./

"""test_transfer.py — the verdict. Run: python -m unittest -v

The tests swap in an UNTRAINED resnet18 (same architecture, no download), so
they run offline and in a couple of seconds. They check how the network is
wired - which parameters are frozen, what the head looks like - not how
accurate it is.
"""

import tempfile
import unittest
from pathlib import Path
from unittest import mock

import torch
from PIL import Image
from torchvision import transforms
from torchvision.models import resnet18

import transfer
from data import make_loaders
from test_data import make_folder


def untrained_resnet18(weights=None):
    return resnet18(weights=None)


class BuildResnetTest(unittest.TestCase):
    def setUp(self):
        patcher = mock.patch.object(transfer, "resnet18", untrained_resnet18)
        patcher.start()
        self.addCleanup(patcher.stop)

    def test_the_head_has_one_output_per_class(self):
        model = transfer.build_resnet(4)
        self.assertEqual(tuple(model.fc.weight.shape), (4, 512))
        out = model.eval()(torch.zeros(2, 3, 64, 64))
        self.assertEqual(tuple(out.shape), (2, 4))

    def test_the_backbone_is_frozen(self):
        model = transfer.build_resnet(4)
        unfrozen = [
            name
            for name, p in model.named_parameters()
            if p.requires_grad and not name.startswith("fc.")
        ]
        self.assertEqual(unfrozen, [], "these backbone parameters are still trainable")

    def test_the_new_head_is_trainable(self):
        model = transfer.build_resnet(4)
        self.assertTrue(model.fc.weight.requires_grad)
        self.assertTrue(model.fc.bias.requires_grad)

    def test_finetuning_unfreezes_layer4_and_nothing_else(self):
        model = transfer.build_resnet(4, finetune_last_block=True)
        for name, p in model.named_parameters():
            should_train = name.startswith("layer4.") or name.startswith("fc.")
            self.assertEqual(p.requires_grad, should_train, name)

    def test_count_trainable_frozen_backbone(self):
        # fc: 512 * 4 weights + 4 biases = 2,052; resnet18 has 11,178,564 with that head
        self.assertEqual(
            transfer.count_trainable(transfer.build_resnet(4)), (2_052, 11_178_564)
        )

    def test_count_trainable_with_layer4(self):
        model = transfer.build_resnet(4, finetune_last_block=True)
        self.assertEqual(transfer.count_trainable(model), (8_395_780, 11_178_564))


class RgbLoaderTest(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.tmp = tempfile.TemporaryDirectory()
        cls.root = make_folder(Path(cls.tmp.name) / "d", {"alpha": 20, "beta": 20})

    @classmethod
    def tearDownClass(cls):
        cls.tmp.cleanup()

    def test_default_loaders_are_still_28px_gray(self):
        train_loader, _, _ = make_loaders(self.root, batch_size=4)
        xb, _ = next(iter(train_loader))
        self.assertEqual(tuple(xb.shape), (4, 1, 28, 28))

    def test_rgb_loaders_give_three_channels_at_the_requested_size(self):
        train_loader, _, _ = make_loaders(self.root, batch_size=4, size=64, rgb=True)
        xb, _ = next(iter(train_loader))
        self.assertEqual(tuple(xb.shape), (4, 3, 64, 64))

    def test_rgb_pixels_use_the_imagenet_statistics(self):
        _, val_loader, _ = make_loaders(self.root, batch_size=4, size=64, rgb=True)
        got, _ = val_loader.dataset[0]
        subset = val_loader.dataset
        path = subset.dataset.samples[subset.indices[0]][0]
        raw = transforms.Compose([transforms.Resize((64, 64)), transforms.ToTensor()])(
            Image.open(path).convert("RGB")
        )
        mean = torch.tensor([0.485, 0.456, 0.406]).reshape(3, 1, 1)
        std = torch.tensor([0.229, 0.224, 0.225]).reshape(3, 1, 1)
        torch.testing.assert_close(got, (raw - mean) / std)


if __name__ == "__main__":
    unittest.main()

count_trainable and RGB pictures — write them yourself

1. count_trainable

2. RGB pictures

The test

When it does not pass

1. `count_trainable`