PyTorch Fine-Tuning
Fine-tuning updates the weights of a pretrained model for your specific task. Unlike feature extraction — which freezes the backbone completely — fine-tuning lets the pretrained layers adapt to your data. Done carefully, it produces higher accuracy than feature extraction, especially when your task differs from the original pretraining task.
Feature Extraction vs Fine-Tuning
Feature Extraction: Freeze backbone → train only new head Fast, few trainable params Best for: small dataset, similar task to pretrained task Fine-Tuning: Start from pretrained weights → unfreeze some or all layers Train at a LOW learning rate to avoid destroying learned features Best for: medium/large dataset, or when task differs from pretrained task
Step 1: Load Pretrained Model and Replace Head
import torch
import torch.nn as nn
import torchvision.models as models
from torchvision.models import ResNet18_Weights
model = models.resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)
# Replace the final layer for your task (e.g. 3 classes)
model.fc = nn.Linear(model.fc.in_features, 3)Step 2: Freeze All Layers, Unfreeze Gradually
A common fine-tuning strategy freezes everything first, trains the head for a few epochs, then unfreezes later layers progressively. This protects learned features during the early unstable phase of head training.
# Phase 1: Freeze backbone, train head only
for name, param in model.named_parameters():
if 'fc' not in name:
param.requires_grad = False
optimizer_phase1 = torch.optim.Adam(
filter(lambda p: p.requires_grad, model.parameters()), lr=1e-3
)
# Train for ~5 epochs with just the head
# Phase 2: Unfreeze entire model, train at much lower lr
for param in model.parameters():
param.requires_grad = True
optimizer_phase2 = torch.optim.Adam([
{'params': model.layer4.parameters(), 'lr': 1e-4},
{'params': model.fc.parameters(), 'lr': 1e-3},
], lr=1e-5) # default lr for all other layers
# Train for ~10–20 more epochs
Discriminative Learning Rates
Early layers contain more general, universal features — they need less updating. Later layers and the new head need more updating. Assigning different learning rates to different parts of the model — lower rates for early layers, higher for later ones — is called discriminative learning rates.
model = models.resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)
model.fc = nn.Linear(512, 3)
optimizer = torch.optim.Adam([
{'params': model.layer1.parameters(), 'lr': 1e-5}, # early, general
{'params': model.layer2.parameters(), 'lr': 2e-5},
{'params': model.layer3.parameters(), 'lr': 5e-5},
{'params': model.layer4.parameters(), 'lr': 1e-4},
{'params': model.fc.parameters(), 'lr': 1e-3}, # new head, fastest
])Diagram: Discriminative Learning Rates on ResNet18
Layer lr Role ──────────────────────────────────────────────── Conv1, BN1 1e-5 Basic edge/color detectors ← barely change Layer1 1e-5 Simple patterns Layer2 2e-5 Shapes and textures Layer3 5e-5 Object parts Layer4 1e-4 High-level semantic features FC (new) 1e-3 Task-specific classification ← change most
Full Fine-Tuning Training Loop
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms, models
from torchvision.models import ResNet18_Weights
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# Data
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225])
])
train_data = datasets.ImageFolder('data/train', transform=transform)
val_data = datasets.ImageFolder('data/val', transform=transform)
train_loader = DataLoader(train_data, batch_size=32, shuffle=True)
val_loader = DataLoader(val_data, batch_size=32)
# Model
model = models.resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)
model.fc = nn.Linear(512, len(train_data.classes))
model = model.to(device)
# All layers trainable at discriminative rates
optimizer = torch.optim.Adam([
{'params': list(model.parameters())[:-2], 'lr': 1e-5},
{'params': model.fc.parameters(), 'lr': 1e-3}
])
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=20)
loss_fn = nn.CrossEntropyLoss()
best_acc = 0.0
for epoch in range(20):
model.train()
for imgs, labels in train_loader:
imgs, labels = imgs.to(device), labels.to(device)
optimizer.zero_grad()
loss = loss_fn(model(imgs), labels)
loss.backward()
optimizer.step()
model.eval()
correct, total = 0, 0
with torch.no_grad():
for imgs, labels in val_loader:
imgs, labels = imgs.to(device), labels.to(device)
preds = model(imgs).argmax(1)
correct += (preds == labels).sum().item()
total += labels.size(0)
acc = correct / total
scheduler.step()
if acc > best_acc:
best_acc = acc
torch.save(model.state_dict(), 'best_finetuned.pth')
print(f"Epoch {epoch+1:2d} Val Acc: {acc:.4f} Best: {best_acc:.4f}")Common Fine-Tuning Mistakes
Mistake 1: Using the same large lr for all layers → Destroys pretrained features in early layers Fix: Use discriminative learning rates Mistake 2: Forgetting to unfreeze layers → Model trains only the head; misses the benefit of fine-tuning Fix: Explicitly unfreeze layers after Phase 1 Mistake 3: Skipping normalization → Model receives input in wrong range, accuracy plummets Fix: Always use ImageNet mean/std normalization Mistake 4: Forgetting model.eval() at validation → BatchNorm and Dropout behave incorrectly Fix: Always switch modes explicitly
Summary
Fine-tuning updates pretrained weights for a new task rather than freezing them. Start by training only the new classification head (Phase 1), then unfreeze more layers and train everything at lower learning rates (Phase 2). Use discriminative learning rates — small rates for early layers, larger for later layers and the new head. Always use the correct input normalization for your pretrained model. Fine-tuning consistently outperforms feature extraction when you have sufficient data (roughly 1,000+ samples per class). It is the dominant approach in modern computer vision and NLP workflows.