PyTorch Transfer Learning

The complete, practical PyTorch implementation of transfer learning — assembling every technique from the Transfer Learning category (freezing, feature extraction, fine-tuning) into runnable code.

Complete Feature-Extraction Example

import torch
import torch.nn as nn
import torchvision.models as models

model = models.resnet50(weights="IMAGENET1K_V2")

for param in model.parameters():
    param.requires_grad = False   # freeze the entire pretrained backbone

model.fc = nn.Linear(model.fc.in_features, num_classes)   # new, trainable head

optimizer = torch.optim.Adam(model.fc.parameters(), lr=1e-3)   # only the new head trains

Complete Partial Fine-Tuning Example

model = models.resnet50(weights="IMAGENET1K_V2")

for name, param in model.named_parameters():
    param.requires_grad = "layer4" in name or "fc" in name   # unfreeze only the last block + head

model.fc = nn.Linear(model.fc.in_features, num_classes)

optimizer = torch.optim.Adam(
    [p for p in model.parameters() if p.requires_grad], lr=1e-4
)

Verifying What's Actually Trainable

trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
total = sum(p.numel() for p in model.parameters())
print(f"Training {trainable:,} of {total:,} parameters ({100*trainable/total:.2f}%)")
# ALWAYS verify this matches your intent -- a common bug is accidentally
# leaving MORE (or fewer) parameters trainable than you meant to

Common Mistakes

• Replacing the final layer before setting requires_grad = False on the rest — this is fine functionally (the new layer's parameters default to requires_grad=True when created), but always double-check the trainable-parameter count afterward to catch any mistakes.
• Passing model.parameters() (all parameters, including frozen ones) to the optimizer instead of filtering to only requires_grad=True parameters — functionally harmless since frozen parameters receive no gradient, but it's cleaner and less error-prone to filter explicitly.

Interview Relevance

Q: "How would you verify, in code, that your transfer learning setup is only training the parameters you intended?" Sum p.numel() across parameters with requires_grad=True and compare against the total parameter count — printing both the trainable count and the percentage of the model actually being trained is a quick, reliable sanity check that catches accidental over- or under-freezing before wasting compute on an incorrectly configured training run.

Practice Question

Write the code to freeze every layer of a pretrained model except its final two layers, assuming you know their parameter names contain "layer4" and "fc".

Related DL Notes

• Freezing Layers
• Fine-Tuning
• Loading PyTorch Models