How do you freeze layers and perform transfer learning / fine-tuning in PyTorch?

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

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

# ── Strategy 1: Feature extraction — freeze ALL pretrained layers
for param in backbone.parameters():
    param.requires_grad = False     # excluded from gradient computation

# Replace the final classification 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)  # NEW layer — requires_grad=True by default

# Only backbone.fc parameters will be updated by the optimizer
optimizer = torch.optim.AdamW(
    filter(lambda p: p.requires_grad, backbone.parameters()),  # only trainable params
    lr=1e-3,
)

# ── Strategy 2: Full fine-tuning with layer-wise (discriminative) learning rates
backbone2 = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V2)
backbone2.fc = nn.Linear(backbone2.fc.in_features, 5)

optimizer2 = torch.optim.AdamW([
    {"params": backbone2.layer1.parameters(), "lr": 1e-5},  # earliest layers — smallest lr
    {"params": backbone2.layer4.parameters(), "lr": 1e-4},  # later layers — bigger lr
    {"params": backbone2.fc.parameters(),     "lr": 1e-3},  # new head — largest lr
])

# ── Verify which parameters are trainable
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:,} ({100*trainable/total:.1f}%)")

# ── Common pattern: train head first, then unfreeze and fine-tune everything
# Phase 1: only train backbone.fc for a few epochs
# Phase 2: unfreeze all layers, train with a small lr to fine-tune end-to-end
for param in backbone.parameters():
    param.requires_grad = True   # unfreeze for phase 2