Transfer learning reuses a model already trained on a large dataset (like ImageNet) as the starting point for a new, often much smaller task โ dramatically reducing the data and compute needed compared to training a CNN from scratch.
Why It Works
A CNN's early layers learn fairly generic, broadly reusable features (edges, textures, simple shapes โ see Feature Map), regardless of the specific final task it was trained for. Only the later layers become highly specialized to the original task's specific classes. This means a backbone pretrained on a huge, diverse dataset already "knows" a great deal that's useful for a new, related task, even one with far less training data of its own.
Two Main Strategies
| Strategy | What Changes | When to Use |
|---|---|---|
| Feature extraction | Freeze the entire pretrained backbone; only train a new final classification layer | Very little new training data available |
| Fine-tuning | Unfreeze some or all of the pretrained backbone, continuing training (often with a small learning rate) | A moderate amount of new training data, and the new task is reasonably related to the original |
Code โ Feature Extraction
import torchvision.models as models
import torch.nn as nn
model = models.resnet18(weights='IMAGENET1K_V1')
for param in model.parameters():
param.requires_grad = False # freeze the entire pretrained backbone
model.fc = nn.Linear(model.fc.in_features, 5) # new head, unfrozen, trained from scratch
# only model.fc's parameters will actually update during training
Code โ Fine-Tuning
model = models.resnet18(weights='IMAGENET1K_V1')
model.fc = nn.Linear(model.fc.in_features, 5)
# leave every parameter trainable, but use a small learning rate to avoid destroying pretrained weights
optimizer = torch.optim.Adam(model.parameters(), lr=1e-5)
Common Mistakes
- Using a normal-sized learning rate when fine-tuning a pretrained model โ this risks catastrophically overwriting valuable pretrained weights in just a few steps; a much smaller learning rate (often 10-100x smaller than training from scratch) is standard.
- Fine-tuning the entire backbone on a very small new dataset โ with too little data, unfreezing everything risks overfitting; feature extraction (or partial fine-tuning of just the later layers) is usually safer in that regime.
Interview Relevance
Q: "Why does transfer learning work โ why are ImageNet-pretrained features useful for an unrelated new task?" Early CNN layers learn fairly generic visual features (edges, textures, simple shapes) that are broadly useful across many visual tasks, not specific to the original dataset's exact classes. Only later layers become highly task-specific โ so reusing the early layers, and only retraining (or lightly fine-tuning) the later ones, transfers much of the useful learned representation to a new task with far less data than training from scratch would require.
Practice Question
You have only 200 labeled images for a new classification task. Would feature extraction or full fine-tuning be the safer starting strategy, and why?