Transfer learning with pretrained models

Reuse what a big model already learned

Training a large network from scratch needs a lot of data and compute. Transfer learning starts from a model pretrained on a large dataset such as ImageNet. Its early layers have learned general features like edges, textures and simple shapes, which are useful for almost any image task.

There are two common modes. In feature extraction, you freeze the pretrained backbone (requires_grad = False on its parameters) and train only a new final layer. It is fast and works well when you have only tens to a few hundred images per class. In fine-tuning, you also unfreeze some or all backbone layers and train them with a small learning rate, which usually gives better accuracy when you have more data.

torchvision provides models with a weights argument, for example models.resnet18(weights=models.ResNet18_Weights.DEFAULT). The weights object's .transforms() gives the exact preprocessing the model expects; using different normalisation silently hurts accuracy. For text and many other domains, the Hugging Face Hub hosts pretrained PyTorch models that follow the same idea.

Feature extraction with ResNet-18

After freezing, only the new head is trainable: 512 x 5 weights plus 5 biases = 2565 parameters. The first run downloads the weights, so it needs internet access.

import torch
from torch import nn
from torchvision import models

weights = models.ResNet18_Weights.DEFAULT
model = models.resnet18(weights=weights)         # downloads on first use
preprocess = weights.transforms()                # matching resize and normalise

for p in model.parameters():
    p.requires_grad = False                      # freeze the backbone

model.fc = nn.Linear(model.fc.in_features, 5)    # new head, trainable by default

trainable = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.Adam(trainable, lr=1e-3)
print(sum(p.numel() for p in trainable))         # 2565

An experienced cook in a new kitchen

A cook trained in North Indian food who joins a South Indian restaurant does not relearn how to hold a knife or control a flame. Those skills (the frozen backbone) transfer. They only learn new recipes (the new head). Given more time, they also refine their technique a little for the new cuisine, which is fine-tuning with a small learning rate.

Mind BatchNorm when fine-tuning

Freezing parameters does not freeze BatchNorm's running statistics; in train() mode they still update from your batches. With very small batches, consider keeping BatchNorm layers in eval mode during fine-tuning, or fine-tune with reasonable batch sizes.

त्वरित जाँच: You have 300 labelled images across 5 classes. What is a sensible first approach?

• Train a ResNet from random initialisation for many epochs
• Fine-tune every layer with a large learning rate
• Use a pretrained model without changing its 1000-class head
• Freeze a pretrained backbone and train a new 5-class head