Fine-Tuning#
Fine-tuning consists of taking a model pretrained on a large dataset, such as ImageNet, and adapting its architecture and parameters to a smaller, task-specific dataset. The approach works because the layers of a convolutional neural network learn features at different levels of abstraction. Early layers typically detect broad visual patterns such as edges, shapes, and textures, while deeper layers capture more specialized, task-dependent representations. By preserving the general features learned by early layers and retraining only selected deep layers on a new dataset, fine-tuning can achieve strong performance with substantially less data and computation than training a model from scratch.
Successful fine-tuning requires careful consideration of several choices, including the pretrained architecture, the number of layers to freeze or unfreeze, the learning rate, and any preprocessing or augmentation applied to the dataset. Poorly chosen settings can reduce performance or damage useful pretrained representations. For example, a newly initialized classification head may initially generate large, noisy gradients, while an excessively high learning rate can rapidly overwrite previously learned features. This can lead to catastrophic forgetting, where the model loses previously learned knowledge.
Important
To avoid catastrophic forgetting, fine-tuning is commonly performed in two stages. First, the pretrained backbone is frozen while the newly added classification head is trained. Next, some of the backbone’s deeper layers are unfrozen and trained using a smaller learning rate, allowing the model to adapt gradually without losing the useful features acquired during pretraining.
This tutorial demonstrates how to build a cats-versus-dogs classifier using MobileNetV3. It explains parameter freezing, the behavior of Batch Normalization layers during fine-tuning, and the use of validation data to guide model development before conducting a single final evaluation on the test set.
Model preparation#
Before we can dive into fine-tuning, we need to understand how to freeze layers and modify the architecture of a pre-trained model. We will illustrate these concepts by working with the MobileNetV3 architecture. As discussed in the previous tutorial, MobileNetV3 consists of three modules.
features- A series of convolutional and pooling layers to learn a representation of the input image.avgpool- A layer that averages the channels of the input 3D tensor to produce a feature vector.classifier- A series of fully-connected layers that predicts the class logits from the feature vector.
Let’s start by instantiating the MobileNetV3 model with randomly initialized weights.
architecture = mobilenet_v3_large()
Trainable layers#
We can inspect the number of trainable tensors in a model by iterating through the model’s parameters and counting how many of them have requires_grad set to True. Alternatively, we can use the torchinfo package to get a summary of the model architecture, including the number of parameters and which ones are trainable. Let’s have an overview of the first two levels of the model. By default, all layers are trainable.
====================================================================================================
Layer (type:depth-idx) Param # Trainable
====================================================================================================
MobileNetV3 -- True
├─Sequential: 1-1 -- True
│ └─Conv2dNormActivation: 2-1 464 True
│ └─InvertedResidual: 2-2 464 True
│ └─InvertedResidual: 2-3 3,440 True
│ └─InvertedResidual: 2-4 4,440 True
│ └─InvertedResidual: 2-5 10,328 True
│ └─InvertedResidual: 2-6 20,992 True
│ └─InvertedResidual: 2-7 20,992 True
│ └─InvertedResidual: 2-8 32,080 True
│ └─InvertedResidual: 2-9 34,760 True
│ └─InvertedResidual: 2-10 31,992 True
│ └─InvertedResidual: 2-11 31,992 True
│ └─InvertedResidual: 2-12 214,424 True
│ └─InvertedResidual: 2-13 386,120 True
│ └─InvertedResidual: 2-14 429,224 True
│ └─InvertedResidual: 2-15 797,360 True
│ └─InvertedResidual: 2-16 797,360 True
│ └─Conv2dNormActivation: 2-17 155,520 True
├─AdaptiveAvgPool2d: 1-2 -- --
├─Sequential: 1-3 -- True
│ └─Linear: 2-18 1,230,080 True
│ └─Hardswish: 2-19 -- --
│ └─Dropout: 2-20 -- --
│ └─Linear: 2-21 1,281,000 True
====================================================================================================
Total params: 5,483,032
Trainable params: 5,483,032
Non-trainable params: 0
====================================================================================================
Freezing layers#
To freeze a layer, we set the requires_grad attribute to False for all parameters in that layer. This prevents the optimizer from updating the weights of the layer during training. We can use the parameters() iterator to access the parameters of a nn.Module and set the requires_grad attribute accordingly. The code below demonstrates how to freeze the convolutional backbone of MobileNetV3.
for params in architecture.features.parameters():
params.requires_grad = False
Let’s regenerate the model summary after freezing the backbone. Only the parameters in the classification head should be trainable now.
====================================================================================================
Layer (type:depth-idx) Param # Trainable
====================================================================================================
MobileNetV3 -- Partial
├─Sequential: 1-1 -- False
│ └─Conv2dNormActivation: 2-1 (464) False
│ └─InvertedResidual: 2-2 (464) False
│ └─InvertedResidual: 2-3 (3,440) False
│ └─InvertedResidual: 2-4 (4,440) False
│ └─InvertedResidual: 2-5 (10,328) False
│ └─InvertedResidual: 2-6 (20,992) False
│ └─InvertedResidual: 2-7 (20,992) False
│ └─InvertedResidual: 2-8 (32,080) False
│ └─InvertedResidual: 2-9 (34,760) False
│ └─InvertedResidual: 2-10 (31,992) False
│ └─InvertedResidual: 2-11 (31,992) False
│ └─InvertedResidual: 2-12 (214,424) False
│ └─InvertedResidual: 2-13 (386,120) False
│ └─InvertedResidual: 2-14 (429,224) False
│ └─InvertedResidual: 2-15 (797,360) False
│ └─InvertedResidual: 2-16 (797,360) False
│ └─Conv2dNormActivation: 2-17 (155,520) False
├─AdaptiveAvgPool2d: 1-2 -- --
├─Sequential: 1-3 -- True
│ └─Linear: 2-18 1,230,080 True
│ └─Hardswish: 2-19 -- --
│ └─Dropout: 2-20 -- --
│ └─Linear: 2-21 1,281,000 True
====================================================================================================
Total params: 5,483,032
Trainable params: 2,511,080
Non-trainable params: 2,971,952
====================================================================================================
Replacing layers#
Fine-tuning often requires modifying the architecture of the pre-trained model to adapt it to the new task. This can involve changing the number of output units in the classification head, adding new layers, or replacing existing layers. There are several ways to modify the architecture of a model in PyTorch. The simplest way is to just replace the layers in the classification head. For example, we can replace the two fully-connected layers of MobileNetV3 with new ones that have a different number of output units.
in_dim = architecture.classifier[0].in_features
architecture.classifier[0] = torch.nn.Linear(in_dim, 500)
architecture.classifier[3] = torch.nn.Linear(500, 1) # binary classification
A more drastic approach is to replace the entire classification head with a new one. This can be done by creating a new nn.Module for the classification head, and replacing the classifier attribute of the model with the new module.
architecture.classifier = torch.nn.Sequential(
torch.nn.Linear(in_dim, 500),
torch.nn.ReLU(),
torch.nn.Linear(500, 1)
)
====================================================================================================
Layer (type:depth-idx) Param # Trainable
====================================================================================================
MobileNetV3 -- Partial
├─Sequential: 1-1 -- False
│ └─Conv2dNormActivation: 2-1 (464) False
│ └─InvertedResidual: 2-2 (464) False
│ └─InvertedResidual: 2-3 (3,440) False
│ └─InvertedResidual: 2-4 (4,440) False
│ └─InvertedResidual: 2-5 (10,328) False
│ └─InvertedResidual: 2-6 (20,992) False
│ └─InvertedResidual: 2-7 (20,992) False
│ └─InvertedResidual: 2-8 (32,080) False
│ └─InvertedResidual: 2-9 (34,760) False
│ └─InvertedResidual: 2-10 (31,992) False
│ └─InvertedResidual: 2-11 (31,992) False
│ └─InvertedResidual: 2-12 (214,424) False
│ └─InvertedResidual: 2-13 (386,120) False
│ └─InvertedResidual: 2-14 (429,224) False
│ └─InvertedResidual: 2-15 (797,360) False
│ └─InvertedResidual: 2-16 (797,360) False
│ └─Conv2dNormActivation: 2-17 (155,520) False
├─AdaptiveAvgPool2d: 1-2 -- --
├─Sequential: 1-3 -- True
│ └─Linear: 2-18 480,500 True
│ └─ReLU: 2-19 -- --
│ └─Linear: 2-20 501 True
====================================================================================================
Total params: 3,452,953
Trainable params: 481,001
Non-trainable params: 2,971,952
====================================================================================================
Wrapping the pretrained architecture#
Replacing the classification head of a pretrained model is often sufficient for many tasks. However, if we want to add custom functionality, such as explicit freeze() and unfreeze() methods, we can create a wrapper class that inherits from nn.Module. The wrapper should preserve the pretrained architecture unless there is a reason to change it.
Important
Batch Normalization needs a deliberate policy during fine-tuning. With small target datasets or small batches, keeping pretrained running statistics fixed is often more stable. With enough representative data and suitably large batches, adapting the statistics and sometimes the affine parameters can improve performance. We use the conservative frozen-statistics policy for this small example and make that choice explicit in the model.
The following helpers change which parameters receive gradients and can keep selected module types in evaluation mode. Remember that requires_grad=False prevents parameter gradients, while eval() controls runtime behavior such as BatchNorm running-statistic updates.
Next, we define a module that preserves both the convolutional backbone and the average pooling layer of MobileNetV3, while replacing the classification head with a new one. The forward() method defines the forward pass of the model, which applies the convolutional backbone, average pooling, flattening, and classification head in sequence. The flattening operation is necessary because the average pooling layer produces a 4D tensor with shape (batch, channels, 1, 1), while the classification head expects a 2D tensor with shape (batch, features). The wrapper module also provides utility methods to freeze and unfreeze the convolutional backbone, as well as to set the Batch Normalization layers to evaluation mode.
class MobileNet(nn.Module):
def __init__(
self,
weights: MobileNet_V3_Large_Weights = None,
freeze_batch_norm: bool = True,
):
super().__init__()
mobilenet = mobilenet_v3_large(weights=weights)
self.backbone = mobilenet.features
self.avgpool = mobilenet.avgpool
self.freeze_batch_norm = freeze_batch_norm
in_features = mobilenet.classifier[0].in_features
self.classifier = nn.Sequential(
nn.Linear(in_features, 500),
nn.ReLU(),
nn.Dropout(0.2),
nn.Linear(500, 1),
)
def forward(self, x):
x = self.backbone(x)
x = self.avgpool(x)
x = torch.flatten(x, 1)
return self.classifier(x).squeeze(1)
#----- Fine-tuning support -----#
def train(self, mode: bool = True):
"""Set training mode while optionally preserving pretrained BN statistics."""
super().train(mode)
if mode and self.freeze_batch_norm:
set_eval_mode(self.backbone, nn.BatchNorm2d)
return self
def freeze(self):
"""Freeze all backbone parameters."""
make_trainable(self.backbone, False)
def unfreeze(self, count: int):
"""Unfreeze the last `count` backbone blocks."""
unfreeze_layers(self.backbone, count)
if self.freeze_batch_norm:
freeze_by_type(self.backbone, nn.BatchNorm2d)
Sanity check#
Let’s check that the model works as expected. First, we create an instance of the model and pass a random batch through it to check the output shape. The final squeezing operation in the forward method returns one logit per image with shape (batch_size,). This is important because the loss function we will use expects a 1D tensor of shape (batch_size,) rather than a 2D tensor of shape (batch_size, 1).
batch = torch.rand(16, 3, 224, 224)
model = MobileNet()
output = model(batch)
assert output.shape == (16,)
Then, we freeze the convolutional backbone and check that only the classification head is trainable.
model.freeze()
====================================================================================================
Layer (type:depth-idx) Param # Trainable
====================================================================================================
MobileNet -- Partial
├─Sequential: 1-1 -- False
│ └─Conv2dNormActivation: 2-1 (464) False
│ └─InvertedResidual: 2-2 (464) False
│ └─InvertedResidual: 2-3 (3,440) False
│ └─InvertedResidual: 2-4 (4,440) False
│ └─InvertedResidual: 2-5 (10,328) False
│ └─InvertedResidual: 2-6 (20,992) False
│ └─InvertedResidual: 2-7 (20,992) False
│ └─InvertedResidual: 2-8 (32,080) False
│ └─InvertedResidual: 2-9 (34,760) False
│ └─InvertedResidual: 2-10 (31,992) False
│ └─InvertedResidual: 2-11 (31,992) False
│ └─InvertedResidual: 2-12 (214,424) False
│ └─InvertedResidual: 2-13 (386,120) False
│ └─InvertedResidual: 2-14 (429,224) False
│ └─InvertedResidual: 2-15 (797,360) False
│ └─InvertedResidual: 2-16 (797,360) False
│ └─Conv2dNormActivation: 2-17 (155,520) False
├─AdaptiveAvgPool2d: 1-2 -- --
├─Sequential: 1-3 -- True
│ └─Linear: 2-18 480,500 True
│ └─ReLU: 2-19 -- --
│ └─Dropout: 2-20 -- --
│ └─Linear: 2-21 501 True
====================================================================================================
Total params: 3,452,953
Trainable params: 481,001
Non-trainable params: 2,971,952
====================================================================================================
Finally, we unfreeze some layers of the convolutional backbone and check that the corresponding parameters are trainable again, except for the BatchNorm layers. In the summary below, we can see that the last three layers of the backbone are marked as partially trainable, while the rest of the backbone remains frozen. Increasing the summary depth would show that the BatchNorm layers are still frozen, as expected.
model.unfreeze(3)
====================================================================================================
Layer (type:depth-idx) Param # Trainable
====================================================================================================
MobileNet -- Partial
├─Sequential: 1-1 -- Partial
│ └─Conv2dNormActivation: 2-1 (464) False
│ └─InvertedResidual: 2-2 (464) False
│ └─InvertedResidual: 2-3 (3,440) False
│ └─InvertedResidual: 2-4 (4,440) False
│ └─InvertedResidual: 2-5 (10,328) False
│ └─InvertedResidual: 2-6 (20,992) False
│ └─InvertedResidual: 2-7 (20,992) False
│ └─InvertedResidual: 2-8 (32,080) False
│ └─InvertedResidual: 2-9 (34,760) False
│ └─InvertedResidual: 2-10 (31,992) False
│ └─InvertedResidual: 2-11 (31,992) False
│ └─InvertedResidual: 2-12 (214,424) False
│ └─InvertedResidual: 2-13 (386,120) False
│ └─InvertedResidual: 2-14 (429,224) False
│ └─InvertedResidual: 2-15 797,360 Partial
│ └─InvertedResidual: 2-16 797,360 Partial
│ └─Conv2dNormActivation: 2-17 155,520 Partial
├─AdaptiveAvgPool2d: 1-2 -- --
├─Sequential: 1-3 -- True
│ └─Linear: 2-18 480,500 True
│ └─ReLU: 2-19 -- --
│ └─Dropout: 2-20 -- --
│ └─Linear: 2-21 501 True
====================================================================================================
Total params: 3,452,953
Trainable params: 2,221,001
Non-trainable params: 1,231,952
====================================================================================================
Fine-tuning workflow#
This is the complete workflow for fine-tuning a pre-trained model on a new task.
Step 1: Select the pretrained model
Select a model pretrained on a large dataset.
Modify the architecture of the pretrained model for the new task.
Step 2: Prepare the dataset
Prepare a dataset for the new task and split it into training, validation, and test sets.
If needed, setup data augmentation on the training set.
Step 3: Train the new head
Freeze the convolutional backbone of the model.
If applicable, set batch normalization layers to evaluation mode.
Create an optimizer over the trainable parameters of the new head.
Train the model on the training set while monitoring the validation set.
Step 4: Adapt part of the backbone
Unfreeze selected layers of the convolutional backbone.
If applicable, keep batch normalization layers frozen and in evaluation mode.
Create a new optimizer over the trainable parameters of the model.
Train the model on the training set with a small learning rate while monitoring the validation set.
Step 5: Evaluate once
After all choices are fixed, evaluate the model on the untouched test set.
Note
Training the head first is a stable default. It prevents an untrained head from immediately sending disruptive gradients into the backbone.
Fine-tuning usually uses a smaller learning rate because pretrained parameters already encode useful structure. More advanced setups use a relatively larger rate for the new head and smaller rates for pretrained layers.
Small datasets or batch sizes may require freezing BatchNorm parameters and running statistics. With larger datasets, it may be beneficial to adapt the statistics and sometimes the affine parameters. Treat this as a hyperparameter and compare policies using validation data.
Step 1: Preparing the model#
We download the weights of a pre-trained MobileNetV3 model and instantiate the model defined earlier.
weights = MobileNet_V3_Large_Weights.DEFAULT
model = MobileNet(weights=weights)
preprocess = weights.transforms()
Step 2: Preparing the data#
We load the cats-vs-dogs dataset and recreate the stratified train/validation/test splits using the preprocessing bundled with the pretrained weights.
See also
The tutorial Dataset from files explains how to download and validate the cats-vs-dogs dataset.
Step 3: Training with frozen backbone#
We import Trainer from the training.py file and BinaryAccuracy from TorchEval.
from training import Trainer
from torcheval.metrics import BinaryAccuracy
Then, we train the model with the convolutional backbone frozen and the classification head unfrozen. Constructing the optimizer after freezing makes its scope explicit and avoids carrying frozen parameters in its parameter groups.
Warning: this cell is computationally expensive without a GPU.
model.freeze()
model = model.to("cuda" if torch.cuda.is_available() else "cpu")
optimizer = optim.Adam(
(parameter for parameter in model.parameters() if parameter.requires_grad),
lr=1e-3,
amsgrad=True,
)
loss_fn = nn.BCEWithLogitsLoss()
epochs = 5
trainer = Trainer()
trainer.set_metrics(accuracy=BinaryAccuracy())
history = trainer.fit(model, train_loader, loss_fn, optimizer, epochs, valid_loader)
===== Training on cuda device =====
Epoch 1/5: 100%|██████████| 274/274 [01:26<00:00, 3.16it/s, accuracy=0.9920, train_loss=0.0421, valid_loss=0.0336]
Epoch 2/5: 100%|██████████| 274/274 [01:21<00:00, 3.36it/s, accuracy=0.9907, train_loss=0.0221, valid_loss=0.0275]
Epoch 3/5: 100%|██████████| 274/274 [01:21<00:00, 3.34it/s, accuracy=0.9835, train_loss=0.0158, valid_loss=0.0429]
Epoch 4/5: 100%|██████████| 274/274 [01:21<00:00, 3.38it/s, accuracy=0.9917, train_loss=0.0120, valid_loss=0.0298]
Epoch 5/5: 100%|██████████| 274/274 [01:21<00:00, 3.36it/s, accuracy=0.9912, train_loss=0.0092, valid_loss=0.0321]
Let’s plot training and validation loss. We use validation behavior to decide whether the frozen-head stage is learning and beginning to overfit.
We evaluate the frozen-backbone model on validation data. This is an intermediate diagnostic, not the final reported test result.
Validation accuracy: 99.12%
Step 4: Fine-tuning selected backbone blocks#
Finally, we unfreeze the last three backbone blocks. Because the set of trainable parameters has changed, we create a new optimizer rather than reusing the head-only optimizer. We use separate parameter groups: a small rate for pretrained backbone parameters and a larger rate for the newer classifier.
Warning: this cell is computationally expensive.
model.unfreeze(3)
optimizer = optim.Adam(
[
{
"params": [
parameter
for parameter in model.backbone.parameters()
if parameter.requires_grad
],
"lr": 1e-5,
},
{"params": model.classifier.parameters(), "lr": 1e-4},
],
amsgrad=True,
)
history = trainer.fit(model, train_loader, loss_fn, optimizer, epochs, valid_loader)
===== Training on cuda device =====
Epoch 1/5: 100%|██████████| 274/274 [01:23<00:00, 3.30it/s, accuracy=0.9920, train_loss=0.0044, valid_loss=0.0328]
Epoch 2/5: 100%|██████████| 274/274 [01:25<00:00, 3.19it/s, accuracy=0.9907, train_loss=0.0017, valid_loss=0.0365]
Epoch 3/5: 100%|██████████| 274/274 [01:24<00:00, 3.23it/s, accuracy=0.9912, train_loss=0.0009, valid_loss=0.0387]
Epoch 4/5: 100%|██████████| 274/274 [01:23<00:00, 3.28it/s, accuracy=0.9920, train_loss=0.0005, valid_loss=0.0399]
Epoch 5/5: 100%|██████████| 274/274 [01:21<00:00, 3.35it/s, accuracy=0.9920, train_loss=0.0003, valid_loss=0.0421]
Let’s plot the second-stage training and validation losses. Comparing them with the frozen stage lets us determine whether adapting the backbone helped.
The validation loss does not show any real improvement, while the validation accuracy improves slightly. This happens because the loss and accuracy metrics answer different questions. Binary accuracy counts how many logits fall on the correct side of the decision threshold; the loss also measures how far and how confidently predictions lie on either side. Accuracy can therefore stay flat while loss changes, or improve while average loss worsens because a few confident mistakes dominate. This is a reason to inspect both metrics, not a reason to disregard increasing validation loss.
After the workflow is fixed, evaluate the fine-tuned model once on the test subset.
Test accuracy: 99.12%
Fine-tuning produces only a small gain compared to feature extraction, since ImageNet features already separate cats and dogs into well-defined clusters. In practice, when the dataset is small or the task is complex, fine-tuning can significantly improve the performance of a pretrained model.
Summary#
In this tutorial, we learned how to fine-tune a pretrained model on a new dataset. We modified the architecture of a pretrained model, prepared the dataset for the new task, and trained the model with a two-stage approach. First, we trained a new classifier on the convolutional backbone frozen. Then, we fine-tuned some layers of the convolutional backbone with a small learning rate. We also discussed the importance of keeping Batch Normalization layers frozen and in evaluation mode during fine-tuning to prevent them from updating their statistics and parameters.