Transfer Learning — Programuoki

Code

build_resnet — fill in two holes

In the last step you saw what has to be done to resnet18. Now you write it as a function — and fill in the two missing pieces yourself.

The algorithm

model ← resnet18 with pretrained weights
for every parameter p in the model:
    switch off its gradients                    # freeze
if finetune_last_block:
    for every parameter p of model.layer4:
        switch gradients on                     # unlock the last block
model.fc ← a new Linear(512 → num_classes)      # new head
return model

The order matters. First you freeze everything, then unlock layer4, and only then replace the head. A freshly created layer is trainable by default — you do nothing for it. If you replaced the head before freezing, the freezing loop would freeze it too.

Your job

Create transfer.py and test_transfer.py (see the right panel).

transfer.py runs but is wrong — two places, HOLE 1 and HOLE 2:

HOLE 1 — freezing one parameter. Right now there is a pass, so nothing is frozen. You need one line with p.requires_grad.
HOLE 2 — the new head. Right now model.fc stays at 1,000 classes; replace it with a single nn.Linear. Take the input size from the old layer: model.fc.in_features.

Unlocking for finetune_last_block is already written.

test_transfer.py is the verdict. It does not download weights: it swaps resnet18 for an untrained network of the same size. It tests not accuracy but the wiring: what is frozen, how big the head is, what is unlocked.

Run it before filling anything in

$ python -m unittest test_transfer
FFF.
======================================================================
FAIL: test_finetuning_unfreezes_layer4_and_nothing_else (...)
AssertionError: True != False : conv1.weight

======================================================================
FAIL: test_the_backbone_is_frozen (...)
AssertionError: Lists differ: ['conv1.weight', 'bn1.weight', 'bn1.bias', ...] != []
...
First list contains 60 additional elements.
======================================================================
FAIL: test_the_head_has_one_output_per_class (...)
AssertionError: Tuples differ: (1000, 512) != (4, 512)
----------------------------------------------------------------------
Ran 4 tests in 0.4s

FAILED (failures=3)

(Shortened.) Read it: 60 parameter tensors of the backbone are still trainable though they should be frozen, and the head has 1,000 outputs, not 4. The one test that passes ("the new head is trainable") passes only because you have not frozen anything yet.

Fill in both holes and run again: you should see OK.

When it does not pass

True != False : conv1.weight — freezing did not reach every parameter. Did you write True instead of False, or p.detach() (which returns a copy and does not change p)?
Tuples differ: (512, 4) != (4, 512) — the nn.Linear arguments are swapped. First is the input (512), second the output (the class count).
False != True : fc.weight (and the head stays at 1,000) — you wrote model.fc.out_features = num_classes. That is only a number on the object; the weight table stays the same. A new layer has to be created. You also see a side effect: the old head had been frozen, so the "new head" is not trainable either.
classifier.weight — classifier is the head of other networks (e.g. VGG); in resnet18 it is called fc. Different networks, different names.

transfer.py CREATE

Create 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

test_transfer.py CREATE

Create 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 unittest
from unittest import mock

import torch
from torchvision.models import resnet18

import transfer


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)


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