What is transfer learning and how do you fine-tune a pretrained model in PyTorch?

Transfer learning reuses a model trained on a large dataset (typically ImageNet for vision, or a large text corpus for NLP) as a starting point for a related task with less data. The pretrained model has already learned general features (edges, textures, shapes for images; grammar, semantics for text) — fine-tuning adapts these features to the target task without needing to learn them from scratch.

Two common strategies: (1) Feature extraction — freeze all pretrained layers and train only a new task-specific head; (2) Full fine-tuning — unfreeze some or all pretrained layers and train end-to-end with a small learning rate to avoid overwriting the useful pretrained representations. A common practical pattern is to first train only the head for a few epochs (so it doesn't start with random gradients corrupting the pretrained backbone), then unfreeze and fine-tune everything together with a smaller lr.

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

# Load pretrained ResNet-50
backbone = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V2)

# --- Strategy 1: Feature extraction ---
# Freeze ALL pretrained parameters
for param in backbone.parameters():
    param.requires_grad = False

# Replace the final FC layer for our task (e.g. 5 classes)
in_features = backbone.fc.in_features  # 2048 for ResNet-50
backbone.fc = nn.Linear(in_features, 5)
# Only backbone.fc.parameters() have requires_grad=True

# --- Strategy 2: Full fine-tuning ---
backbone2 = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V2)
backbone2.fc = nn.Linear(backbone2.fc.in_features, 5)
# Use layer-wise lr: smaller lr for early layers
optimizer = torch.optim.AdamW([
    {'params': backbone2.layer1.parameters(), 'lr': 1e-5},
    {'params': backbone2.layer4.parameters(), 'lr': 1e-4},
    {'params': backbone2.fc.parameters(),    'lr': 1e-3},
], weight_decay=1e-2)

# Verify which parameters will be updated
trainable = sum(p.numel() for p in backbone.parameters() if p.requires_grad)
total     = sum(p.numel() for p in backbone.parameters())
print(f'Trainable: {trainable:,} / Total: {total:,}')