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

Transfer learning reuses a model pretrained on a large dataset and adapts it to a new task. Freezing layers (setting requires_grad=False) prevents their weights from updating during backpropagation — useful when you want to keep pretrained features fixed and only train a new task-specific head.

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

Invest now in Acorns!!! 🚀
Join Acorns and get your $5 bonus!

Acorns is a micro-investing app that automatically invests your "spare change" from daily purchases into diversified, expert-built portfolios of ETFs. It is designed for beginners, allowing you to start investing with as little as $5. The service automates saving and investing. Disclosure: I may receive a referral bonus.

Invest now!!! Get Free equity stock (US, UK only)!

Use Robinhood app to invest in stocks. It is safe and secure. Use the Referral link to claim your free stock when you sign up!.

The Robinhood app makes it easy to trade stocks, crypto and more.

Webull! Receive free stock by signing up using the link: Webull signup.