Learn AI Create the Future — with Pranay
Main site

Building Models with nn.Module

Transfer learning: swapping the head and freezing a backbone

Take a model trained on millions of images, freeze its hard-won feature detectors, replace its final layer with one sized for your classes, and train only that — strong results from small data.

Read these first

On this page 5
  1. Why it exists
  2. How it works
  3. A real example you have seen
  4. Remember this
  5. What to learn next

One lesson, three depths. Pick the one that fits you today — you can switch any time.

Beginner — No maths. Plain English.

Transfer learning means starting from a model someone already trained, keeping most of it, and re-training only a small final part on your own data.

You want a great chef for your new restaurant. One option: raise a child, teach them knife skills, spices, timing — twenty years and a fortune. The other: hire an experienced chef and spend two weeks teaching them your menu.

The chef's fundamentals — heat, salt, texture — transfer to any menu. Only the dishes are new. Transfer learning hires the chef.

Why it exists

Training a top image model from nothing takes millions of labelled images and serious compute. You have 500 photos of three kinds of mango leaf disease, taken on a phone.

Here is the saving grace: the early parts of an image model learn things that have nothing to do with its original task. Edges, textures, spots, shapes — the visual fundamentals. Those transfer. Only the final step — "which of MY categories is this?" — is specific to you.

So the recipe: keep the trained backbone (the fundamentals), remove the final layer — the head — and bolt on a new head sized for your classes. Freeze the backbone, meaning its numbers stay untouched, and train only the new head. Suddenly 500 photos is enough.

How it works

pretrained model:  [ backbone: edges → textures → shapes ] → [ head: 1000 old classes ]
                                                                       |
                                                              remove the head
                                                                       |
your model:        [ backbone: FROZEN ]                    → [ new head: your 3 classes ]
                          |                                            |
                    stays as trained                          the only part that trains

A real example you have seen

Nearly every "AI detects crop disease / skin condition / defect" app began as a model trained on a giant general photo collection, with its head swapped for the specific task. This is the most-used deep learning technique in industry — most real projects start from a pretrained model, not from zero.

Remember this

  • Visual fundamentals transfer; only the final decision is task-specific.
  • Freeze the backbone, swap the head, train the small new part.
  • This is how small datasets get big-model results.

What to learn next

  • Fine-tuning — the same idea grown up, across vision and language.
  • LoRA — freezing everything and training tiny inserted pieces instead.
  • Writing your own Dataset — feeding your 500 photos to the model you now know how to build.

Developer — Code and libraries.

Setup

bash
pip install torch torchvision

Written and tested against torch 2.5 and torchvision 0.20 on CPU.

Freeze, swap, count

The code uses weights=None — random weights, no download — so you can run it instantly and see the mechanics. For real transfer learning, use the pretrained weights (noted below; a one-time ~45 MB download).

freeze_and_swap.py
import torch
from torch import nn
from torchvision import models

# weights=None builds the architecture with random weights: no download.
# For real transfer learning use weights=models.ResNet18_Weights.DEFAULT (~45 MB, once).
model = models.resnet18(weights=None)

print("original head:", model.fc)

for p in model.parameters():          # freeze everything first
    p.requires_grad_(False)

model.fc = nn.Linear(model.fc.in_features, 3)   # new layers arrive unfrozen

trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
frozen = sum(p.numel() for p in model.parameters() if not p.requires_grad)
print(f"trainable: {trainable:,}")
print(f"frozen:    {frozen:,}")

params = [p for p in model.parameters() if p.requires_grad]
opt = torch.optim.AdamW(params, lr=1e-3)
print("tensors the optimizer sees:", len(params))
Output
original head: Linear(in_features=512, out_features=1000, bias=True)
trainable: 1,539
frozen:    11,176,512
tensors the optimizer sees: 2

Eleven million numbers frozen; 1,539 in training. That ratio is the entire economics of transfer learning.

The walkthrough

Freeze first, swap second. requires_grad_(False) walks every parameter and switches off gradient tracking. The new nn.Linear is created after, so it arrives with requires_grad=True by default. Reverse the order and you freeze your new head too — a classic.

model.fc is found by reading, not guessing. Print the model. torchvision's ResNets call the head fc; its VGGs use classifier[6]; EfficientNets use classifier[1]. One print(model) beats any memorised table.

model.fc.in_features reads the incoming width (512 here) off the old head so the new one fits exactly. Never hard-code it.

Give the optimizer only the trainable parameters. Passing everything works — frozen tensors get no gradients — but Adam-family optimizers then keep state for eleven million dead weights. Filtering is cleaner and lighter, and the printed 2 (weight and bias of the new head) confirms exactly what is training. Frozen parameters remain registered and saved, as the registration lesson explains — freezing changes training, not membership.

The BatchNorm trap, one more time

Freezing does not stop BatchNorm's running statistics from updating — that is controlled by train/eval mode, a different switch. Left in train mode, your frozen backbone quietly rewrites its statistics to your new data; sometimes helpful, sometimes destabilising, never intended if you did not choose it. The explicit recipe:

python
model.train()                 # head trains normally
for m in model.modules():
    if isinstance(m, nn.BatchNorm2d):
        m.eval()              # frozen backbone's notebooks stay shut

Run this every epoch after model.train(), since train() resets every submodule. Details in BatchNorm.

When to unfreeze more

Head-only training works when your images resemble the pretraining photos. The further your data drifts — X-rays, satellite tiles, microscope slides — the more layers deserve unfreezing. The standard escalation: train the head; then unfreeze the last block with a learning rate 10x smaller; unfreeze more only if validation says so. Small learning rates protect pretrained weights from being trampled early. The full spectrum, up to LLM-style adaptation, lives in fine-tuning.

Common mistakes

Wrong preprocessing. Pretrained weights assume the exact normalisation used in pretraining. With torchvision, use models.ResNet18_Weights.DEFAULT.transforms() and feed it your images — it carries the correct resize and normalisation with the weights.

Training metrics computed in the wrong mode. See the BatchNorm section above; symptoms are validation scores that wobble with batch size.

A learning rate sized for scratch training. 1e-3 on an unfrozen pretrained backbone destroys it in a few steps. Unfrozen pretrained layers want 1e-5 to 1e-4.

Forgetting the head is random. For the first epochs, gradients from a random head flow into everything unfrozen. This is the argument for head-first training before any unfreezing.

Try it yourself

Unfreeze model.layer4 as well and rebuild the optimizer with two parameter groups: head at 1e-3, layer4 at 1e-4. Print the trainable count and check it against the head-only run.

What to learn next

  • Fine-tuning — the same idea grown up, across vision and language.
  • LoRA — freezing everything and training tiny inserted pieces instead.
  • Writing your own Dataset — feeding your 500 photos to the model you now know how to build.

Researcher — Mathematics and papers.

What transfers, measured

Yosinski et al. (2014), How transferable are features in deep neural networks?, quantified the layer-wise story: early layers converge to Gabor-like filters and colour blobs regardless of task ("general"), later layers specialise, and transferability degrades smoothly with depth. Two further findings outlast the AlexNet era: fine-tuning transferred features beats freezing when target data suffices, and splitting co-adapted layers hurts — fragile co-adaptation is a cost distinct from specificity. Kornblith et al. (2019), Do Better ImageNet Models Transfer Better?, adds the practical regression: ImageNet accuracy predicts transfer accuracy with high correlation across 12 datasets, with the correlation weakening on fine-grained tasks.

Linear probing versus fine-tuning

Head-only training is linear probing: the backbone becomes a fixed feature map $\phi(x)$ and the problem is convex in the head — immune to feature destruction but bounded by frozen features. Kumar et al. (2022), Fine-Tuning can Distort Pretrained Features and Underperform Out-of-Distribution, formalises the tension: full fine-tuning from a random head lets early large-gradient updates distort $\phi$, and the LP-FT recipe — probe first, then fine-tune from the probed head — dominates both extremes, especially out-of-distribution. This is the theoretical backing for the escalation ladder in the developer block.

Discriminative learning rates (per-depth LR decay) and gradual unfreezing arrived with ULMFiT (Howard and Ruder, 2018), transplanted from vision folklore into NLP and back.

Parameter-efficient variants

Freezing-and-swapping generalises to inserting small trainable structures in a frozen network: adapters (Houlsby et al., 2019), and low-rank updates — LoRA (Hu et al., 2021), where $\Delta W = BA$ with rank $r \ll \min(d_{in}, d_{out})$, merging into $W$ at inference with zero latency cost. The head-swap of this lesson is the $r{=}\text{full}$, last-layer-only corner of that design space. Trade-offs are measured in trainable-parameter fraction versus accuracy gap; on many benchmarks sub-1% trainable parameters recover within noise of full fine-tuning.

Failure modes with citations

Negative transfer — pretraining that hurts — appears when source and target distributions conflict (Wang et al., 2019 survey). He, Girshick, Dollár (2019), Rethinking ImageNet Pre-training, showed random initialisation matches pretraining on COCO given 10x the training schedule: pretraining buys sample efficiency and wall-clock, not a higher ceiling, when target data is plentiful. Batch statistics mismatch (the BatchNorm trap above) is quantified in detection literature, where frozen-BN backbones are standard practice precisely to avoid small-batch statistic corruption.

Reading

  • Yosinski et al. (2014); Kornblith et al. (2019) — what transfers, and how well.
  • Kumar et al. (2022) — LP-FT and feature distortion.
  • Hu et al. (2021), LoRA: Low-Rank Adaptation of Large Language Models.

What to learn next

  • Fine-tuning — the same idea grown up, across vision and language.
  • LoRA — freezing everything and training tiny inserted pieces instead.
  • Writing your own Dataset — feeding your 500 photos to the model you now know how to build.