Transfer Learning, Model Fine-Tuning & LoRA

1. The Two Transfer Learning Paradigms

2. Using torchvision's Modern Weights API

In modern PyTorch torchvision, the legacy pretrained=True boolean flag is deprecated in favor of explicit enum weights that bundle normalization statistics and preprocessing transforms:

import torch
import torch.nn as nn
from torchvision.models import resnet50, ResNet50_Weights

# 1. Fetch official ImageNet pretrained weights and auto-transforms
weights = ResNet50_Weights.DEFAULT
preprocess_transforms = weights.transforms()

print("Required Input Preprocessing:")
print(preprocess_transforms)
# Output: Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])

# 2. Instantiate model with weights
model = resnet50(weights=weights)

3. Freezing Layers & Swapping the Classification Head

To freeze the backbone, iterate through model.parameters() and disable autograd tracking. Then, replace the final fc layer with a new projection matching your dataset's class count:

# 1. Freeze all backbone feature extractor weights
for param in model.parameters():
    param.requires_grad = False

# 2. Inspect the existing final classification layer
# For ResNet, the final layer is named 'fc' with in_features=2048
num_features = model.fc.in_features

# 3. Replace with a custom classifier head for 5 custom medical classes
# Newly created layers automatically have requires_grad=True!
model.fc = nn.Sequential(
    nn.Linear(num_features, 256),
    nn.ReLU(),
    nn.Dropout(p=0.4),
    nn.Linear(256, 5)  # 5 custom target classes
)

# 4. Verify trainable parameters
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
total = sum(p.numel() for p in model.parameters())
print(f"Trainable Parameters: {trainable:,} / {total:,} ({trainable/total*100:.2f}%)")
# Only ~500k weights train instead of all 25.5 million!

4. Discriminative Learning Rates in Optimizers

When unfreezing the network for stage-two fine-tuning, applying the same learning rate across all layers will destroy the fragile generic low-level edge filters learned by earlier layers. Pass separate parameter groups with differential learning rates to your optimizer:

# Unfreeze top layers for stage-two fine-tuning
for param in model.layer4.parameters():
    param.requires_grad = True

# Differential Parameter Groups
optimizer = torch.optim.AdamW([
    # Backbone deep layers: ultra-small learning rate to fine-tune without catastrophic forgetting
    {'params': model.layer4.parameters(), 'lr': 1e-5, 'weight_decay': 0.01},
    # Newly initialized head: standard learning rate to learn task mappings rapidly
    {'params': model.fc.parameters(), 'lr': 1e-3, 'weight_decay': 0.01}
])

5. Modern Frontier: Parameter-Efficient Fine-Tuning (PEFT / LoRA)

For Massive Language Models (e.g., LLaMA, Mistral, Gemma with 7B+ parameters), full fine-tuning requires prohibitive VRAM. Low-Rank Adaptation (LoRA) freezes the original weight matrix $W_0$ and injects two low-rank matrices $A$ and $B$:

Where $r \ll d$ (e.g. rank $r = 8$ or $16$). This reduces trainable parameters by 99.9%, allowing multi-billion parameter models to be fine-tuned on a single consumer GPU!