Feature Extraction: Freezing the Backbone

Transfer learning has two main strategies. In feature extraction, the pre-trained backbone is completely frozen — its weights do not change during training. Only the new classification head, which you add on top, learns from your data. In fine-tuning, the entire network or at least some backbone layers are also updated.

Feature extraction is the right choice when your dataset is small (less than a few thousand images) or when your images are similar to ImageNet (natural photographs of everyday objects). It is much faster since gradients do not flow through the backbone, and it avoids destroying carefully learned features with noisy updates from too little data.

In PyTorch, each parameter tensor has a requires_grad attribute. Setting it to False prevents gradient computation for that tensor, effectively freezing it. The simplest way to freeze all backbone parameters is to iterate over model.parameters() and set requires_grad = False, then replace the classification head (which starts with new random weights, so requires_grad=True by default).

This is efficient: PyTorch's autograd skips frozen parameters during the backward pass, reducing memory usage and speeding up training significantly compared to fine-tuning the whole network.

ImageNet pre-trained models learn a hierarchy of features: early layers detect low-level edges and colours, middle layers detect textures and parts, and later layers detect abstract objects. These features are general visual features that transfer broadly across image domains.

When your task involves natural images — flowers, animals, medical scans, satellite photos — these pre-learned features are far more informative than anything a randomly initialised network could extract from a small dataset. Feature extraction leverages this by treating the frozen backbone as a fixed feature transformer and only training a small linear classifier on top of the extracted features.

Once features are frozen, training becomes fast. The optimiser only updates the new head's weights. You can either train a simple nn.Linear layer in PyTorch, or even pass the pre-extracted feature vectors to scikit-learn's LogisticRegression or SVC — both approaches work well for linear classification on top of rich pre-trained features.

Using a single linear layer is equivalent to training a logistic regression on the extracted features. For more complex tasks or when your classes require non-linear decision boundaries, you can use a small multi-layer head with ReLU activations and dropout between the frozen backbone and the output.

During feature extraction, image preprocessing must match the transforms used when training the original model. For ImageNet pre-trained models this means: resize to 256, centre-crop to 224×224, convert to tensor, and normalise with ImageNet mean and std.

When you use weights.transforms(), PyTorch automatically provides the correct preprocessing pipeline associated with those specific weights. This eliminates a common source of subtle bugs where you accidentally use wrong normalisation constants, which can reduce transfer learning accuracy by several percentage points.

Feature extraction is dramatically faster than full fine-tuning. When all backbone parameters have requires_grad=False, PyTorch does not compute gradients through them during loss.backward(), saving memory and computation proportional to the number of frozen layers.

For ResNet-50 (25M parameters, 2048 frozen features, only ~2K parameters in the head), a training epoch on a 5000-image dataset completes in seconds on CPU, compared to minutes for full fine-tuning. Pre-extracting features offline and then training only the head as a standard sklearn classifier is even faster, since the backbone forward pass only runs once per image.

The approach is the same for EfficientNet: freeze all parameters, then replace the final classification layer. For EfficientNet the classifier is a Sequential block accessed via model.classifier, not model.fc. Always check the model architecture to find the correct attribute name.

EfficientNet-B0's classifier expects 1280-dimensional input features. When writing code that works with multiple architectures, introspect the final layer with in_features instead of hard-coding the dimension, making your code more reusable across different backbone choices.

For ViT, the classification head is accessed via model.heads. The feature dimension from ViT-B/16's backbone is 768 (the transformer embedding dimension). After freezing all parameters, replace model.heads with a new nn.Linear(768, num_classes).

ViT benefits less from feature extraction than CNNs because its attention layers are more task-specific — they learn to attend to different image regions than what your task requires. If you use a ViT backbone, at minimum unfreeze the last few transformer encoder blocks for best performance. However, for a quick experiment, full feature extraction is still a valid starting point.

After training with frozen features, evaluate on the validation set using standard classification metrics. A well-tuned feature extractor often achieves 90%+ accuracy on small custom datasets (500-5000 images per class) where training from scratch would overfit completely.

If accuracy is lower than expected, consider: (1) adding data augmentation during training, (2) using a dropout layer before the linear head, (3) using a slightly more expressive head (two linear layers with ReLU), or (4) moving to partial fine-tuning by unfreezing the last residual block. Monitor both train and validation accuracy to distinguish overfitting from underfitting.

Feature extraction works best when the source domain (ImageNet: natural photographs) and target domain are similar. When your data is very different — X-ray images, satellite imagery, microscopy, infrared thermal images — the low-level features still transfer (edge detectors are universal), but higher-level ImageNet features may be less useful.

Signs that feature extraction is underperforming: validation accuracy plateaus well below expectations, or train accuracy is much higher than validation (indicating the head is overfitting to too few features). In these cases, move to fine-tuning: unfreeze some backbone layers and train with a much lower learning rate to gradually adapt the pre-learned features to your domain.

Test your understanding of feature extraction with frozen backbones from this lesson.

In this lesson you learned: feature extraction freezes all backbone parameters and trains only a new classification head, giving fast training and strong results when your data resembles ImageNet, freezing is done by setting requires_grad=False on backbone parameters before replacing the final layer, and pre-extracting features offline is even faster by running the backbone once and caching the feature vectors. Next up we learn fine-tuning — how to carefully unfreeze backbone layers with low learning rates for even better accuracy.