Table of Contents
What Freezing and Fine-Tuning Mean
When you use a pretrained vision model, most of the useful visual features are already stored in its weights. Freezing means you prevent some parameters from being updated during training. Fine-tuning means you allow parameters to update so the model adapts to your dataset. In practice you often replace the model head for your task, then decide how much of the backbone to freeze or fine-tune.
A common pattern is to freeze the entire backbone at first and train only a newly added classifier head. If results plateau, you progressively unfreeze later layers and fine-tune them with a smaller learning rate.
Replacing the Classification Head
Pretrained models are typically trained on ImageNet with 1000 classes, so the last layer does not match your number of classes. You usually keep the backbone and swap the final classifier layer.
In torchvision, different architectures expose the classifier differently. For example, ResNet uses model.fc, many MobileNet and EfficientNet variants use model.classifier, and VGG uses model.classifier as a sequential block. The key idea is always the same, you create a new nn.Linear with output size equal to your classes.
Important rule: you must replace the head before constructing the optimizer, otherwise the optimizer might not include the new parameters.
Freezing Parameters Correctly
In PyTorch, a parameter is updated by gradient descent only if two things happen, it has requires_grad=True, and it is included in the optimizer parameter list. Freezing is usually done by setting requires_grad=False for selected parameters.
A typical workflow is to freeze everything, then unfreeze just the head.
import torch
import torch.nn as nn
from torchvision import models
num_classes = 5
model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
for p in model.parameters():
p.requires_grad = False
in_features = model.fc.in_features
model.fc = nn.Linear(in_features, num_classes)
for p in model.fc.parameters():
p.requires_grad = TrueAt this point, only the head will produce gradients.
Important rule: freezing with requires_grad=False stops gradient computation for those parameters, but you still must ensure your optimizer only receives the parameters you intend to train.
Building the Optimizer for Only Trainable Parameters
To train only the head, pass only trainable parameters to the optimizer. A robust pattern is filtering by requires_grad.
optimizer = torch.optim.Adam(
(p for p in model.parameters() if p.requires_grad),
lr=1e-3
)This avoids accidentally training frozen layers and automatically includes new layers you added.
Unfreezing in Stages
If your dataset is moderately sized and different from ImageNet, training only the head may not be enough. A common next step is to unfreeze the last block or the last few layers of the backbone and fine-tune them. With ResNet, later blocks are often layer4, then layer3, and so on.
for p in model.layer4.parameters():
p.requires_grad = TrueAfter changing which parameters require gradients, you should recreate the optimizer so it includes the newly trainable parameters.
Important rule: when you unfreeze additional layers, recreate the optimizer, otherwise those parameters will not be updated because the optimizer does not know about them.
Using Different Learning Rates for Head and Backbone
Newly initialized layers like your head usually need a larger learning rate than pretrained layers, which you want to change more gently. You can set parameter groups in the optimizer.
head_params = list(model.fc.parameters())
backbone_params = [p for n, p in model.named_parameters()
if not n.startswith("fc.") and p.requires_grad]
optimizer = torch.optim.Adam([
{"params": backbone_params, "lr": 1e-5},
{"params": head_params, "lr": 1e-3},
])This is one of the most practical tricks for stable fine-tuning.
Important rule: use a smaller learning rate for pretrained layers than for the newly initialized head, often 10 times to 100 times smaller.
BatchNorm and the Difference Between Freezing Weights and Freezing Behavior
Layers like BatchNorm have parameters and also internal running statistics that update during training when the model is in train() mode. Even if you freeze BatchNorm parameters with requires_grad=False, the running mean and variance can still change during training, which can hurt when your dataset is small.
A common approach is to keep the model in training mode overall but set BatchNorm layers to evaluation mode so their running statistics do not update.
def set_bn_eval(m):
if isinstance(m, nn.BatchNorm2d):
m.eval()
model.apply(set_bn_eval)This is not always required, but it is a frequent fix when fine-tuning becomes unstable on small datasets.
Important rule: requires_grad=False freezes gradients, but it does not automatically stop BatchNorm running statistics updates. Control this with train() and eval() behavior.
When to Freeze vs Fine-Tune
Freezing most of the network is often best when your dataset is small, similar to the pretrained domain, or you want fast training and lower risk of overfitting. Fine-tuning more layers is often best when your dataset is larger, significantly different from ImageNet, or you need extra accuracy and can afford more compute.
A practical heuristic is to start with head only, then progressively unfreeze from the end of the backbone, watching validation performance.
A Minimal Fine-Tuning Checklist
Before you start training, confirm that the head output matches the number of classes, that only intended parameters have requires_grad=True, that the optimizer was created after freezing and head replacement, and that you are using a smaller learning rate for pretrained layers if you fine-tune them. During training, if validation accuracy drops sharply when you unfreeze, reduce the backbone learning rate and consider freezing BatchNorm behavior.