Transfer learning reuses a model that already learned useful visual features and adapts it to a new task with less labeled data than training from scratch. In PyTorch computer vision work, the usual outcome is a classifier, detector, or segmenter whose early layers keep generic edge, texture, color, and shape detectors while later layers are replaced or adjusted for the new label space. Fine-tuning is the controlled part of that adaptation: you decide which parameters stay frozen, which parameters receive gradients, and how aggressively each group is updated.
In this course, transfer learning sits after the lessons on tensors, modules, optimizers, and convolutional vision models because it combines all of them. A pretrained network is still an ordinary nn.Module. The difference is that its parameters begin from a useful point in weight space instead of random initialization, and your training loop must respect which parts of the network are acting as a fixed feature extractor and which parts are being changed.
Purpose and Outcome
The purpose is not to blindly import a famous architecture. The purpose is to reduce data requirements, training time, and overfitting risk when your target images share visual structure with the pretraining data. A medical image classifier, warehouse defect detector, wildlife camera classifier, or plant disease model can often start with a pretrained vision backbone and learn a smaller amount of task-specific information in a new head.
By the end of a good transfer learning run you should have evidence that the adapted model beats a simple baseline on held-out data, behaves sensibly on important slices, and has a reproducible recipe: preprocessing, backbone choice, head shape, frozen layers, optimizer parameter groups, learning rates, checkpoint selection, and evaluation protocol.
How It Works Internally
A deep vision model transforms an image through successive feature maps. Early convolutional blocks usually learn local patterns such as edges and simple textures. Middle blocks compose those patterns into parts and repeated motifs. Deeper blocks become more tied to the source task’s categories. Transfer learning exploits this uneven generality. You keep the layers likely to remain useful and replace the layers whose output vocabulary no longer matches your task.
In PyTorch, freezing is implemented through each Parameter‘s requires_grad flag. When requires_grad=False, autograd does not store gradients for that parameter and optimizers will not update it if the parameter is excluded from their parameter groups. The forward pass still runs through the frozen layers, so the model can still compute features. Only learning is disabled for those weights.
The classifier head is the common replacement point. For an image classification model, a pretrained backbone may produce a feature vector and an old final layer may output one thousand source classes. If your dataset has five classes, the final layer must output five logits. That head is initialized randomly, so it usually receives a larger learning rate than the pretrained body. If accuracy plateaus after the head learns, you may unfreeze later backbone blocks and continue with a smaller learning rate so useful source features are nudged rather than overwritten.
API Anatomy
The essential PyTorch pieces are small but important. The model is an nn.Module; parameters are returned by model.parameters() or named subsets such as model.classifier.parameters(); trainability is controlled by parameter.requires_grad; update rules live in an optimizer; and mode-dependent layers are controlled by model.train() and model.eval(). Batch normalization and dropout make the mode call especially important. A frozen parameter can still sit inside a module whose batch normalization running statistics change during training mode, so freezing weights and controlling module mode are related but not identical choices.
Preprocessing is part of the transferred model’s contract. The input size, color channel order, normalization mean and standard deviation, and value range should match the assumptions used by the pretrained weights. If the source model expected normalized RGB tensors and you feed unnormalized grayscale-expanded images, transfer can fail even though the tensor shape is accepted.
Example 1: Replace a Classifier Head
This first example isolates the head replacement idea with a small tensor instead of a full image model. A backbone would emit three features per image, and the new classifier maps those features to two target classes.
import torch
from torch import nn
features = torch.tensor([[0.2, 0.8, 0.4], [0.9, 0.1, 0.3]])
classifier = nn.Linear(3, 2)
logits = classifier(features)
print(logits.shape)The deterministic expected behavior is the printed shape: torch.Size([2, 2]). There are two input items and two output logits per item. In a real classifier, those logits would be passed to cross_entropy with integer class targets. The important design point is that the final dimension of the head must equal the number of target classes, not the number of classes in the source dataset.
Example 2: Freeze the Backbone
The second example turns a small module into a feature extractor. The backbone still computes features, but its weights are not trainable. Only the new head contributes trainable parameters.
import torch
from torch import nn
class TinyBackbone(nn.Module):
def __init__(self):
super().__init__()
self.features = nn.Sequential(nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 6), nn.ReLU())
def forward(self, x):
return self.features(x)
backbone = TinyBackbone()
for parameter in backbone.parameters():
parameter.requires_grad = False
head = nn.Linear(6, 3)
model = nn.Sequential(backbone, head)
print(sum(p.numel() for p in model.parameters() if p.requires_grad))The expected output is 21, because a Linear(6, 3) head has eighteen weights and three bias values. This is the typical first stage for a small dataset: train only the randomly initialized head while preserving the pretrained representation. It is faster, uses less optimizer state, and reduces overfitting, but it can underfit if the new domain is visually far from the source domain.
Example 3: Staged Fine-Tuning with Parameter Groups
After the head has learned a reasonable mapping, you can unfreeze a later block and fine-tune it cautiously. Parameter groups let you assign a small learning rate to pretrained layers and a larger learning rate to the new classifier.
import torch
from torch import nn
class SmallTransferModel(nn.Module):
def __init__(self):
super().__init__()
self.features = nn.Sequential(nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 6), nn.ReLU())
self.classifier = nn.Linear(6, 3)
def forward(self, x):
return self.classifier(self.features(x))
model = SmallTransferModel()
for parameter in model.features[0].parameters():
parameter.requires_grad = False
optimizer = torch.optim.SGD([
{"params": model.features[2].parameters(), "lr": 1e-4},
{"params": model.classifier.parameters(), "lr": 1e-2},
], momentum=0.9)
x = torch.randn(5, 4)
y = torch.tensor([0, 1, 2, 1, 0])
loss = nn.functional.cross_entropy(model(x), y)
optimizer.zero_grad(set_to_none=True)
loss.backward()
optimizer.step()
print(round(float(loss.detach()), 4))The final printed loss is numeric but not deterministic because the inputs and initialization are random. The behavior that matters is structural: the first feature layer remains frozen, the later feature layer receives small updates, and the classifier receives larger updates. This staged approach often improves accuracy when the target task differs from the source task but still shares low-level visual patterns.
Example 4: Verify Inference Output
Transfer learning is useful only if the adapted model produces well-formed predictions. This example checks that inference returns one probability distribution per item.
import torch
from torch import nn
model = nn.Sequential(nn.Linear(4, 6), nn.ReLU(), nn.Linear(6, 2))
batch = torch.randn(3, 4)
with torch.no_grad():
probabilities = torch.softmax(model(batch), dim=1)
print(probabilities.shape)
print(torch.allclose(probabilities.sum(dim=1), torch.ones(3)))The expected output is torch.Size([3, 2]) followed by True. The first line confirms three examples and two classes. The second confirms that each row sums to one after softmax. In production code you would also map class indices back to stable label names and test that the preprocessing used at inference matches training.
Design Choices and Trade-Offs
The first choice is source model selection. A larger backbone can provide stronger features but increases latency, memory use, and risk of overfitting when fine-tuned on a tiny dataset. A smaller backbone may be easier to deploy and iterate on. Match the architecture to image size, deployment device, and the amount of labeled data, not just leaderboard reputation.
The second choice is freeze depth. Training only the head is stable and cheap. Unfreezing the last block gives the model room to adapt semantic features. Unfreezing the whole network gives maximum flexibility but can destroy useful pretrained weights if the learning rate is too high or the dataset is small. This forgetting of transferred knowledge is a common reason that a fine-tuned model performs worse than a frozen-backbone baseline.
The third choice is optimization schedule. Use lower learning rates for pretrained layers, consider warmup for newly unfrozen blocks, and checkpoint by validation metric rather than final epoch. If classes are imbalanced, pair transfer learning with sampling, weighting, or a metric such as per-class recall instead of relying only on average accuracy.
Failure Modes and Troubleshooting
Symptom: training loss falls but validation accuracy stays flat. A likely cause is that the head is learning shortcuts or the split leaks near-duplicate images. Diagnose by inspecting per-class confusion, nearest neighbors across train and validation folders, and a baseline trained on shuffled labels. Correct it by rebuilding the split by subject, scene, time, or source device, and by adding augmentation that matches realistic variation.
Symptom: loss becomes unstable immediately after unfreezing. The common cause is a learning rate that is acceptable for a random head but too large for pretrained layers. Diagnose by logging gradient norms for the unfrozen block and comparing validation loss before and after the unfreeze epoch. Correct it with separate parameter groups, a lower backbone learning rate, and a staged schedule that unfreezes only later blocks first.
Symptom: validation works, but live inference is poor. The usual cause is preprocessing drift: image resizing, cropping, color conversion, normalization, or label mapping differs between training and serving. Diagnose by saving a preprocessed tensor from both paths for the same image and comparing shape, min, max, mean, and class-index mapping. Correct it by packaging preprocessing with the model or sharing one transform definition across training and inference.
Symptom: frozen models still change between epochs. Batch normalization running statistics may still update when modules remain in training mode. Diagnose by comparing running_mean tensors before and after an epoch. Correct it by setting frozen batch normalization modules to evaluation mode when appropriate, while keeping the trainable head in training mode.
Reliability and Performance Implications
Transfer learning reduces training cost, but it can hide dataset mismatch. A model pretrained on everyday images may encode features and biases that are irrelevant or harmful for specialized domains. Reliability comes from slice-based evaluation: check rare classes, acquisition devices, lighting conditions, geographic regions, and edge cases separately. Security also matters when loading weights. Treat model checkpoints as code-like artifacts: obtain them from trusted sources, record hashes where possible, and avoid loading untrusted serialized objects in privileged environments.
For inference performance, most cost remains in the backbone, not the new head. Freezing speeds training but does not automatically speed inference. To improve serving latency, consider smaller architectures, static input sizes, batching, quantization supported by your deployment target, or caching features only when images are immutable and the cache key includes preprocessing version and backbone version.
Hands-On Lab
Prerequisites: a Python environment with PyTorch installed, a small image dataset arranged by class or a synthetic tensor dataset for a dry run, and enough disk space for checkpoints. If using a torchvision pretrained model, install torchvision that is compatible with your PyTorch installation.
- Create train, validation, and test splits with no near-duplicate images crossing splits.
- Define transforms for training and validation. Keep validation deterministic and record image size and normalization.
- Load a pretrained backbone, replace the classifier head so its output size equals your class count, and freeze all backbone parameters.
- Train the head for a few epochs. Save the checkpoint with the best validation metric.
- Unfreeze the last backbone block, rebuild the optimizer with a smaller learning rate for that block, and continue training from the best head-only checkpoint.
- Evaluate once on the untouched test split. Report confusion matrix, per-class recall, and the selected checkpoint name.
- Run one inference sample end to end and verify the predicted index maps to the intended class label.
Verification: confirm that the head output shape equals [batch_size, num_classes], frozen parameters have grad is None after backward, validation metrics improve over a head-only or from-scratch baseline, and test metrics are computed only after model selection. Cleanup: delete exploratory checkpoints that were not selected, keep the final checkpoint with its preprocessing configuration, and roll back by restoring the previous checkpoint and label map.
Assessment Exercises
- You have 600 labeled images across three classes and a validation curve that improves only after unfreezing the last block. Explain why this can happen and which learning rates you would compare.
- A model trained with frozen weights changes validation predictions between epochs. List two PyTorch mechanisms that could cause this and how you would inspect each one.
- Design an experiment to decide between a small backbone trained with more augmentation and a large backbone with only the head trained.
- Your live service receives correct tensor shapes but wrong classes. Write a diagnostic plan that separates preprocessing mismatch from label-index mismatch.
- When would training from scratch be more defensible than transfer learning for a computer vision task?
Summary
Transfer learning in PyTorch is the practice of starting from useful pretrained visual representations, replacing task-specific outputs, and controlling which parameters adapt. Fine-tuning is effective when it is staged, measured, and tied to the target dataset’s visual distance from the source task. The core mechanics are ordinary PyTorch mechanics: modules, parameters, gradients, optimizer groups, modes, transforms, and checkpoints. Use frozen-backbone baselines, cautious unfreezing, slice-based evaluation, and reproducible preprocessing to turn a convenient shortcut into a defensible computer vision workflow.
