Transfer Learning

By the end of this lesson you will be able to take a network that someone else trained on a huge dataset and adapt it to your own small problem. You will load pretrained weights with torchvision, freeze them, attach a new output layer, and fine-tune with care. This is how most real image and text projects start.

Training a deep network from scratch needs a lot of labelled data and computing power. A college project on crop diseases might have 2,000 leaf photos and one laptop. Transfer learning makes that enough.

The idea: reuse what was already learned

Transfer learning means starting from a model trained on one large task (the source) and reusing it for your own task (the target). A model trained on ImageNet, with over a million photos in 1,000 classes, has already learned to see edges, textures, shapes and object parts. Those skills are useful for almost any photo task.

Think of an experienced cook from a busy Delhi restaurant who joins a restaurant in Kochi. Knife skills, timing and hygiene all carry over on day one. Only the new recipes need learning. Nobody sends the cook back to learn how to hold a knife.

Inside a CNN, as you saw in CNN Architectures: LeNet to ResNet:

• early layers detect general things like edges and colour blobs, useful for any images;
• middle layers detect textures and parts;
• the last layers combine them into the source task's classes, such as "goldfish" or "tiger shark", which you probably do not need.

So you keep the body of the network, called the backbone, and replace its last layer, called the head.

Two strategies

Freezing a layer means its weights do not change during training. In PyTorch you do that by setting requires_grad = False on its parameters, so autograd computes no gradients for them and the optimiser leaves them alone.

A common recipe combines both: first train only the head for a few epochs, then unfreeze the backbone and continue with a much smaller learning rate.

Step 1: load pretrained weights

torchvision provides many pretrained models. You choose the weights with an enum such as ResNet18_Weights.DEFAULT, which means "the best weights available for this model". The weights object also carries the class names and the exact preprocessing the model was trained with. The first call downloads the weights file (about 45 MB) and caches it. This is plain code with its real output.

import torch
from torchvision.models import resnet18, ResNet18_Weights

weights = ResNet18_Weights.DEFAULT
model = resnet18(weights=weights)
print("weights:", weights)
print("ImageNet classes:", len(weights.meta["categories"]), "| first five:", weights.meta["categories"][:5])
print(weights.transforms())

weights: ResNet18_Weights.IMAGENET1K_V1
ImageNet classes: 1000 | first five: ['tench', 'goldfish', 'great white shark', 'tiger shark', 'hammerhead']
ImageClassification(
    crop_size=[224]
    resize_size=[256]
    mean=[0.485, 0.456, 0.406]
    std=[0.229, 0.224, 0.225]
    interpolation=InterpolationMode.BILINEAR
)

The transform tells you what the model expects: resize to 256, crop the centre 224 × 224, and normalise each colour channel with ImageNet's mean and standard deviation. Your images must be prepared the same way. A pretrained model fed differently scaled images performs much worse, and nothing warns you.

Step 2: freeze the backbone and replace the head

ResNet-18's head is model.fc, a linear layer from 512 features to 1,000 classes. For a 4-class task, say four kinds of leaf disease on a farm near Nashik, we freeze everything and then put in a new nn.Linear(512, 4). A newly created layer has requires_grad=True, so only it will train.

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

model = resnet18(weights=ResNet18_Weights.DEFAULT)
for p in model.parameters():
    p.requires_grad = False                         # freeze the whole backbone
model.fc = nn.Linear(model.fc.in_features, 4)       # new head for 4 classes

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

optimizer = torch.optim.AdamW(trainable, lr=1e-3)   # give the optimiser only the head
x = torch.randn(2, 3, 224, 224)
print("output:", tuple(model(x).shape))

total parameters:     11,178,564
trainable parameters: 2,052
output: (2, 4)

Out of over 11 million parameters, only about two thousand will be trained. That is why feature extraction works with a few hundred images per class and trains quickly even on a CPU: most of the network only runs forward.

Step 3: the same mechanics in NumPy

Here is the idea in miniature, small enough to run in your browser. A "pretrained" first layer is frozen, and only the new head is trained. After training, the frozen weights are exactly as they were.

import numpy as np
from sklearn.datasets import load_digits

X, y = load_digits(return_X_y=True)
X = X / 16.0
rng = np.random.default_rng(0)

W_backbone = rng.normal(0, 0.3, (64, 32))      # imagine these came from a pretrained model
b_backbone = np.zeros(32)
W_head, b_head = np.zeros((32, 10)), np.zeros(10)   # new head, trained from zero
frozen_copy = W_backbone.copy()

features = np.maximum(0, X @ W_backbone + b_backbone)   # backbone only runs forward
for step in range(300):
    Z = features @ W_head + b_head
    P = np.exp(Z - Z.max(axis=1, keepdims=True))
    P /= P.sum(axis=1, keepdims=True)
    d = P.copy()
    d[np.arange(len(y)), y] -= 1
    d /= len(y)
    W_head -= 0.5 * features.T @ d                  # update the head only
    b_head -= 0.5 * d.sum(axis=0)

print("trainable:", W_head.size + b_head.size, "| frozen:", W_backbone.size + b_backbone.size)
print("backbone unchanged:", np.array_equal(W_backbone, frozen_copy))
print(f"training accuracy with the head alone: {np.mean(P.argmax(axis=1) == y):.3f}")

trainable: 330 | frozen: 2080
backbone unchanged: True
training accuracy with the head alone: 0.843

Because the backbone never changes, its features can even be computed once and stored. Then each epoch only trains the small head, which is very fast. In this toy, the "backbone" is random rather than truly pretrained, yet a head on top of it already does well; a real pretrained backbone gives far more useful features.

Step 4: fine-tuning with care

When you unfreeze the backbone, its weights are already good. Large updates would destroy what they learned. So fine-tuning uses:

1. a much smaller learning rate, often 10 to 100 times smaller than for training from scratch;
2. different learning rates for different parts: tiny for early layers, larger for later layers and the head;
3. gradual unfreezing: first the head, then the last block, then more.

PyTorch supports point 2 with parameter groups: give the optimiser a list of dictionaries, each with its own lr.

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

model = resnet18(weights=ResNet18_Weights.DEFAULT)
model.fc = nn.Linear(512, 4)
for name, p in model.named_parameters():
    p.requires_grad = name.startswith(("layer4", "fc"))   # unfreeze the last block and the head

optimizer = torch.optim.AdamW([
    {"params": model.layer4.parameters(), "lr": 1e-4},     # pretrained block: small steps
    {"params": model.fc.parameters(), "lr": 1e-3},         # new head: larger steps
], weight_decay=0.01)
for group in optimizer.param_groups:
    print(f"lr {group['lr']}: {sum(p.numel() for p in group['params']):,} parameters")

lr 0.0001: 8,393,728 parameters
lr 0.001: 2,052 parameters

One subtle point: batch normalisation layers keep updating their running statistics in model.train() mode, even when their weights are frozen. With small or unusual datasets, some people keep the frozen part in evaluation mode for this reason. Try both if fine-tuning behaves strangely.

Beyond images

The same idea drives modern natural language processing. Models such as BERT are pretrained on huge amounts of text and then fine-tuned for tasks like classifying support tickets or detecting spam in Hindi and English messages. Libraries such as Hugging Face Transformers make loading these models and their matching tokenisers a few lines of code. Large language models take this further: the Prompting vs RAG vs Fine-tuning lesson compares the options.

When transfer helps, and when it does not

Transfer works best when the source and target are related: everyday photos transfer well to product photos or leaf photos. It helps less for very different data, such as X-rays, satellite images or spectrograms, though starting from pretrained weights is still usually better than random ones. If you have plenty of data that looks nothing like the source, fine-tuning the whole network, or even training from scratch, may win. Always compare against a baseline on your validation set.

Check the licence of any pretrained weights before using them in a product. Some are for research use only.

Common mistakes

• Using different preprocessing from the pretrained model. Use weights.transforms() or the same mean, standard deviation and size.
• Giving the optimiser the frozen parameters. Pass only the parameters with requires_grad=True, or use parameter groups.
• Fine-tuning with a large learning rate. It wipes out the pretrained knowledge in a few steps.
• Grey images into a colour model. Repeat the single channel three times (for example v2.Grayscale(num_output_channels=3)).
• Forgetting that a new head starts random. Train it first with the backbone frozen, then unfreeze.

Key points

• Transfer learning reuses a model pretrained on a large dataset for a new, smaller task.
• Keep the backbone, replace the head with a new layer sized for your classes.
• Feature extraction trains only the head; fine-tuning also updates backbone layers with a small learning rate.
• Freeze layers with requires_grad = False and give the optimiser only trainable parameters.
• Prepare inputs exactly as the pretrained model expects, using weights.transforms().

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

model = resnet18(weights=ResNet18_Weights.DEFAULT)
model.fc = nn.Linear(model.fc.in_features, 7)
for name, p in model.named_parameters():
    p.requires_grad = name.startswith(("layer3", "layer4", "fc"))

print(f"trainable: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}")
optimizer = torch.optim.AdamW([
    {"params": model.layer3.parameters(), "lr": 5e-5},
    {"params": model.layer4.parameters(), "lr": 1e-4},
    {"params": model.fc.parameters(), "lr": 1e-3},
])
print([group["lr"] for group in optimizer.param_groups])

trainable: 10,497,031
[5e-05, 0.0001, 0.001]