Transfer Learning

import torch
import torch.nn as nn
import torchvision
from torchvision import models, transforms, datasets
from torch.utils.data import DataLoader
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR

# 1) Load a pretrained model
weights = models.ResNet50_Weights.IMAGENET1K_V2
base = models.resnet50(weights=weights)
print(base)                                            # final layer is Linear(2048, 1000)

# Use the recommended preprocessing
preprocess = weights.transforms()

# 2) Swap the head for your task
NUM_CLASSES = 10
base.fc = nn.Linear(base.fc.in_features, NUM_CLASSES)   # 1000 -> NUM_CLASSES

# 3) Strategy A: feature extraction (freeze everything but the head)
for p in base.parameters():
    p.requires_grad = False
for p in base.fc.parameters():
    p.requires_grad = True

# Quick to train, modest accuracy. Best when your dataset is tiny + similar to pretraining domain.

# 4) Strategy B: full fine-tune (every layer trains)
for p in base.parameters():
    p.requires_grad = True

# Best accuracy when you have ~ 5k+ labelled samples. Needs lower LR than from-scratch training.

# 5) Strategy C: layer-wise LR — gold standard
# Higher LR for the new head, lower for early layers.
param_groups = [
    {'params': list(base.fc.parameters()),       'lr': 1e-3},
    {'params': list(base.layer4.parameters()),   'lr': 1e-4},
    {'params': list(base.layer3.parameters()),   'lr': 5e-5},
    {'params': list(base.layer1.parameters()) + list(base.layer2.parameters()), 'lr': 1e-5},
    {'params': list(base.conv1.parameters()) + list(base.bn1.parameters()),     'lr': 1e-5},
]
opt = AdamW(param_groups, weight_decay=1e-4)

# Or — discriminative learning rates per parameter group (fastai-style).

# 6) Training loop
import os, time
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
base = base.to(device)

train_ds = datasets.ImageFolder('data/train', transform=preprocess)
val_ds   = datasets.ImageFolder('data/val',   transform=preprocess)
train_dl = DataLoader(train_ds, batch_size=64, shuffle=True,  num_workers=4, pin_memory=True)
val_dl   = DataLoader(val_ds,   batch_size=64, shuffle=False, num_workers=4, pin_memory=True)

criterion = nn.CrossEntropyLoss()
scheduler = CosineAnnealingLR(opt, T_max=10)

for epoch in range(10):
    base.train()
    for x, y in train_dl:
        x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)
        opt.zero_grad(set_to_none=True)
        loss = criterion(base(x), y)
        loss.backward()
        opt.step()
    scheduler.step()

    base.eval()
    correct = 0; total = 0
    with torch.no_grad():
        for x, y in val_dl:
            preds = base(x.to(device)).argmax(dim=1).cpu()
            correct += (preds == y).sum().item(); total += y.size(0)
    print(f'epoch {epoch+1} val acc {correct/total:.3f}')

# 7) Mixed precision — 2-3x faster, lower memory
scaler = torch.cuda.amp.GradScaler()
for x, y in train_dl:
    x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)
    opt.zero_grad(set_to_none=True)
    with torch.cuda.amp.autocast(dtype=torch.float16):
        loss = criterion(base(x), y)
    scaler.scale(loss).backward()
    scaler.step(opt)
    scaler.update()

# 8) Gradual unfreezing (Howard & Ruder, ULMFiT)
# Phase 1: train only the head (5 epochs)
# Phase 2: unfreeze last block, continue at lower LR (3 epochs)
# Phase 3: unfreeze everything, even lower LR (3-5 epochs)
# Often produces the best generalisation, especially with little data.

# 9) Other vision backbones
models.efficientnet_v2_s(weights=models.EfficientNet_V2_S_Weights.IMAGENET1K_V1)
models.convnext_tiny(weights=models.ConvNeXt_Tiny_Weights.IMAGENET1K_V1)
models.vit_b_16(weights=models.ViT_B_16_Weights.IMAGENET1K_SWAG_E2E_V1)
models.swin_v2_t(weights=models.Swin_V2_T_Weights.IMAGENET1K_V1)

# Pick by speed/memory budget vs accuracy on your benchmark.

# 10) Hugging Face — NLP transfer learning
from transformers import AutoModelForSequenceClassification, AutoTokenizer, TrainingArguments, Trainer

model = AutoModelForSequenceClassification.from_pretrained('roberta-base', num_labels=4)
tok   = AutoTokenizer.from_pretrained('roberta-base')

args = TrainingArguments(
    output_dir='out',
    num_train_epochs=3,
    per_device_train_batch_size=16,
    learning_rate=2e-5,
    weight_decay=0.01,
    warmup_ratio=0.06,
    eval_strategy='epoch',
    save_strategy='epoch',
    load_best_model_at_end=True,
    fp16=True,
)
trainer = Trainer(model=model, args=args, train_dataset=train_hf, eval_dataset=val_hf)
trainer.train()

# 11) PEFT — Parameter-Efficient Fine-Tuning (LoRA, adapters)
# When the model is huge (Llama, Mistral), fine-tuning all weights is expensive.
# LoRA trains tiny low-rank adapter matrices, leaves base weights frozen.
from peft import LoraConfig, get_peft_model, TaskType

config = LoraConfig(
    task_type=TaskType.SEQ_CLS,
    r=8, lora_alpha=16, lora_dropout=0.05,
    target_modules=['query', 'value'],
)
model = get_peft_model(model, config)
model.print_trainable_parameters()        # 0.1-1% of full model parameters

# Trains 100x faster, fits a 7B model on a single 24GB GPU.

# 12) Data efficiency — augmentation makes transfer go further
train_transform = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ColorJitter(0.2, 0.2, 0.2),
    transforms.RandAugment(num_ops=2, magnitude=9),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# Use MixUp / CutMix for an extra 1-2 points on tough datasets:
from torchvision.transforms import v2
cutmix_or_mixup = v2.RandomChoice([v2.CutMix(num_classes=NUM_CLASSES), v2.MixUp(num_classes=NUM_CLASSES)])

# 13) When transfer doesn't help
# • Domain gap is huge (ImageNet → medical X-ray): freezing all but the head underperforms a smaller bespoke model
# • Your dataset is millions of labels: a from-scratch model can beat a frozen pretrained one
# • You need very small + fast inference: distil a smaller model from your fine-tuned one
# • The pretrained model was trained on data you legally / ethically can't build on

# 14) Save + load checkpoints (state_dict is the recommended format)
torch.save(base.state_dict(), 'fine_tuned.pt')
base.load_state_dict(torch.load('fine_tuned.pt', map_location='cpu'))

# 15) Quantisation post-fine-tune
import torch.ao.quantization as q
base.eval()
base_q = q.quantize_dynamic(base, {nn.Linear}, dtype=torch.qint8)
torch.save(base_q.state_dict(), 'fine_tuned_qint8.pt')
# 4x smaller; same accuracy on most tasks; faster CPU inference.

# 16) Common bugs
# • Forgot to swap the head — model still outputs 1000 classes
# • Used the SAME LR for head and backbone — backbone overfits to small data
# • Wrong preprocessing pipeline — mean/std mismatch hurts accuracy 5-10%
# • Did not call model.eval() at inference — BatchNorm/Dropout active, wrong predictions
# • Trained too long with a tiny dataset — overfits; use early stopping, dropout, weight decay
# • Loaded weights without strict=False after swapping head — KeyError on fc / classifier
# • Mixed-precision with optimizer.zero_grad() but forgot set_to_none=True — memory grows
# • Forgot to freeze BatchNorm running stats during full fine-tune — set to .eval() per layer to keep stats from drifting on small batches

import torchvision.models as M
from torch import nn
base = M.resnet50(weights=M.ResNet50_Weights.IMAGENET1K_V2)
for p in base.parameters(): p.requires_grad = False
base.fc = nn.Linear(base.fc.in_features, num_classes)