Transfer Learning — Fine-Tuning Pretrained Vision Models

A civil engineer who has spent 10 years building roads knows structural principles, material properties, load calculations — universal engineering knowledge. To now build bridges, they do not retrain their entire education. They learn bridge-specific design on top of existing expertise. That is transfer learning.

Early layers of a pretrained CNN have learned to detect edges, curves, and textures from 1.2 million ImageNet images. That knowledge transfers perfectly to detecting fabric defects, product damage, or medical anomalies. Only the final task-specific layers need to learn from your small dataset.

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

# ── Strategy 1: Feature extraction — freeze everything except head ─────
def build_feature_extractor(n_classes: int, backbone: str = 'resnet50'):
    model = getattr(models, backbone)(weights='IMAGENET1K_V2')

    # Freeze ALL backbone parameters
    for param in model.parameters():
        param.requires_grad = False

    # Replace the head — only this trains
    in_features = model.fc.in_features
    model.fc = nn.Sequential(
        nn.Dropout(0.3),
        nn.Linear(in_features, n_classes),
    )
    # model.fc is unfrozen by default

    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total     = sum(p.numel() for p in model.parameters())
    print(f"Feature extractor: {trainable:,} / {total:,} trainable ({trainable/total*100:.2f}%)")
    return model

# ── Strategy 2: Fine-tuning — freeze early, unfreeze late ─────────────
def build_finetune_model(n_classes: int, backbone: str = 'resnet50',
                          unfreeze_layers: list = ['layer3', 'layer4', 'fc']):
    model = getattr(models, backbone)(weights='IMAGENET1K_V2')

    # Freeze everything first
    for param in model.parameters():
        param.requires_grad = False

    # Selectively unfreeze later layers
    for layer_name in unfreeze_layers:
        layer = getattr(model, layer_name, None)
        if layer:
            for param in layer.parameters():
                param.requires_grad = True

    # Replace head
    in_features = model.fc.in_features
    model.fc    = nn.Sequential(
        nn.Dropout(0.4),
        nn.Linear(in_features, 512),
        nn.ReLU(),
        nn.Dropout(0.3),
        nn.Linear(512, n_classes),
    )

    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total     = sum(p.numel() for p in model.parameters())
    print(f"Fine-tune model:   {trainable:,} / {total:,} trainable ({trainable/total*100:.2f}%)")
    return model

# ── Compare both ──────────────────────────────────────────────────────
fe_model  = build_feature_extractor(n_classes=6)
ft_model  = build_finetune_model(n_classes=6, unfreeze_layers=['layer3','layer4','fc'])

# Verify: which layers are trainable in fine-tune model?
print("
Trainable layers in fine-tune model:")
for name, param in ft_model.named_parameters():
    if param.requires_grad:
        print(f"  ✓ {name}")

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.models as models
import torchvision.transforms as T
import numpy as np
from PIL import Image
from torch.utils.data import Dataset, DataLoader
from sklearn.metrics import accuracy_score
import copy, warnings
warnings.filterwarnings('ignore')

torch.manual_seed(42)

# ── Build model ───────────────────────────────────────────────────────
model = models.resnet50(weights='IMAGENET1K_V2')
model.fc = nn.Sequential(
    nn.Dropout(0.4),
    nn.Linear(model.fc.in_features, 6),
)

# ── Differential learning rates — param groups ────────────────────────
BASE_LR = 1e-3

param_groups = [
    {'params': model.layer1.parameters(), 'lr': BASE_LR * 0.01},
    {'params': model.layer2.parameters(), 'lr': BASE_LR * 0.01},
    {'params': model.layer3.parameters(), 'lr': BASE_LR * 0.1},
    {'params': model.layer4.parameters(), 'lr': BASE_LR * 1.0},
    {'params': model.fc.parameters(),     'lr': BASE_LR * 10},
    # Freeze layer0 (conv1 + bn1) — very early edge detectors, never touch
    # {'params': model.layer0.parameters(), 'lr': 0},   # effectively frozen
]

# Freeze layer0 explicitly
for param in model.conv1.parameters(): param.requires_grad = False
for param in model.bn1.parameters():   param.requires_grad = False

optimizer = optim.AdamW(param_groups, weight_decay=0.01)

# Verify lr per group
print("Differential learning rates:")
group_names = ['layer1', 'layer2', 'layer3', 'layer4', 'fc']
for name, group in zip(group_names, optimizer.param_groups):
    n_params = sum(p.numel() for p in group['params'])
    print(f"  {name:<8}: lr={group['lr']:.0e}  params={n_params:,}")

# ── Gradual unfreezing — start frozen, unfreeze layer by layer ────────
# Alternative to differential lr — especially good for very small datasets

def gradual_unfreeze_schedule(model, epoch):
    """
    Epoch 1-5:   Only head trains
    Epoch 6-10:  Unfreeze layer4
    Epoch 11-15: Unfreeze layer3
    Epoch 16+:   Unfreeze layer2
    """
    # Start: freeze all backbone
    for param in model.parameters():
        param.requires_grad = False
    for param in model.fc.parameters():
        param.requires_grad = True

    if epoch >= 6:
        for param in model.layer4.parameters():
            param.requires_grad = True
    if epoch >= 11:
        for param in model.layer3.parameters():
            param.requires_grad = True
    if epoch >= 16:
        for param in model.layer2.parameters():
            param.requires_grad = True

    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total     = sum(p.numel() for p in model.parameters())
    return trainable, total

print("
Gradual unfreezing schedule:")
test_model = models.resnet50(weights=None)
test_model.fc = nn.Linear(test_model.fc.in_features, 6)
for epoch in [1, 6, 11, 16, 20]:
    tr, tot = gradual_unfreeze_schedule(test_model, epoch)
    print(f"  Epoch {epoch:2d}: {tr:>8,} / {tot:,} trainable ({tr/tot*100:.1f}%)")

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

# ── Load and adapt any backbone in a unified way ──────────────────────
def build_classifier(backbone_name: str, n_classes: int,
                      dropout: float = 0.3) -> nn.Module:
    """
    Universal function to build a transfer learning classifier
    from any torchvision backbone.
    """
    weights_map = {
        'resnet18':      'IMAGENET1K_V1',
        'resnet50':      'IMAGENET1K_V2',
        'resnet101':     'IMAGENET1K_V2',
        'efficientnet_b0': 'IMAGENET1K_V1',
        'efficientnet_b3': 'IMAGENET1K_V1',
        'vit_b_16':      'IMAGENET1K_V1',
        'convnext_small': 'IMAGENET1K_V1',
    }
    weights = weights_map.get(backbone_name, 'IMAGENET1K_V1')
    model   = getattr(models, backbone_name)(weights=weights)

    # Get the in_features of the final classification layer
    # Different backbones use different attribute names
    if hasattr(model, 'fc'):           # ResNet family
        in_features = model.fc.in_features
        model.fc    = nn.Sequential(
            nn.Dropout(dropout),
            nn.Linear(in_features, n_classes),
        )
    elif hasattr(model, 'classifier'): # EfficientNet, ConvNeXt
        if isinstance(model.classifier, nn.Sequential):
            in_features = model.classifier[-1].in_features
        else:
            in_features = model.classifier.in_features
        model.classifier = nn.Sequential(
            nn.Dropout(dropout),
            nn.Linear(in_features, n_classes),
        )
    elif hasattr(model, 'heads'):      # ViT
        in_features   = model.heads.head.in_features
        model.heads.head = nn.Sequential(
            nn.Dropout(dropout),
            nn.Linear(in_features, n_classes),
        )

    return model

# ── Benchmark latency for each backbone ──────────────────────────────
N_CLASSES = 6
dummy     = torch.randn(1, 3, 224, 224)

print(f"{'Backbone':<20} {'Params':>8} {'Latency (ms)':>14}")
print("─" * 46)

for name in ['resnet18', 'resnet50', 'efficientnet_b0', 'efficientnet_b3']:
    model = build_classifier(name, N_CLASSES)
    model.eval()

    params = sum(p.numel() for p in model.parameters()) / 1e6

    # Warmup
    with torch.no_grad():
        for _ in range(5): _ = model(dummy)

    # Time
    start = time.time()
    with torch.no_grad():
        for _ in range(50): _ = model(dummy)
    ms = (time.time() - start) / 50 * 1000

    print(f"  {name:<18}  {params:>6.1f}M  {ms:>12.2f}ms")

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.models as models
import torchvision.transforms as T
import numpy as np
from PIL import Image
from torch.utils.data import Dataset, DataLoader
from sklearn.metrics import accuracy_score, classification_report
import copy, warnings
warnings.filterwarnings('ignore')

torch.manual_seed(42)
np.random.seed(42)

CATEGORIES = ['kurta', 'saree', 'jeans', 'sneakers', 'watch', 'handbag']
N_CLASSES   = len(CATEGORIES)
IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD  = [0.229, 0.224, 0.225]

# ── Dataset ───────────────────────────────────────────────────────────
class ShopifyDataset(Dataset):
    def __init__(self, n=500, split='train'):
        np.random.seed(42 if split == 'train' else 7)
        self.labels = np.random.randint(0, N_CLASSES, n)
        self.bias   = np.array([
            [.8,.3,.2],[.2,.7,.5],[.2,.3,.8],
            [.7,.7,.2],[.5,.5,.5],[.6,.2,.6],
        ])
        if split == 'train':
            self.transform = T.Compose([
                T.RandomResizedCrop(224, scale=(0.7, 1.0)),
                T.RandomHorizontalFlip(),
                T.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.2),
                T.ToTensor(),
                T.Normalize(IMAGENET_MEAN, IMAGENET_STD),
            ])
        else:
            self.transform = T.Compose([
                T.Resize(256), T.CenterCrop(224),
                T.ToTensor(),
                T.Normalize(IMAGENET_MEAN, IMAGENET_STD),
            ])

    def __len__(self): return len(self.labels)

    def __getitem__(self, i):
        b   = self.bias[self.labels[i]]
        img = np.random.randn(64, 64, 3) * 0.15
        for c in range(3): img[:, :, c] += b[c]
        img = np.clip(img, 0, 1)
        img = Image.fromarray((img * 255).astype(np.uint8))
        return self.transform(img), self.labels[i]

train_ds = ShopifyDataset(400, 'train')
val_ds   = ShopifyDataset(100, 'val')
train_ld = DataLoader(train_ds, 32, shuffle=True,  drop_last=True)
val_ld   = DataLoader(val_ds,   32, shuffle=False)

# ── Model with differential lr ─────────────────────────────────────────
model   = models.resnet50(weights='IMAGENET1K_V2')
in_feat = model.fc.in_features
model.fc = nn.Sequential(nn.Dropout(0.4), nn.Linear(in_feat, N_CLASSES))

# Freeze layer1 and layer2
for param in model.layer1.parameters(): param.requires_grad = False
for param in model.layer2.parameters(): param.requires_grad = False

BASE_LR  = 1e-3
optimizer = optim.AdamW([
    {'params': model.layer3.parameters(), 'lr': BASE_LR * 0.1},
    {'params': model.layer4.parameters(), 'lr': BASE_LR * 1.0},
    {'params': model.fc.parameters(),     'lr': BASE_LR * 10},
], weight_decay=0.01)

criterion = nn.CrossEntropyLoss(label_smoothing=0.1)
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=20)

# ── Training loop with early stopping ────────────────────────────────
best_acc, best_wts, patience = 0.0, None, 0

print(f"Fine-tuning ResNet50 on Shopify product categories:")
print(f"{'Epoch':>6} {'Train loss':>12} {'Val acc':>10}")
print("─" * 32)

for epoch in range(1, 21):
    model.train()
    total_loss = 0
    for imgs, labels in train_ld:
        optimizer.zero_grad()
        loss = criterion(model(imgs), labels)
        loss.backward()
        nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        total_loss += loss.item()
    scheduler.step()

    model.eval()
    all_preds, all_labels = [], []
    with torch.no_grad():
        for imgs, labels in val_ld:
            preds = model(imgs).argmax(1)
            all_preds.extend(preds.numpy())
            all_labels.extend(labels.numpy())

    val_acc = accuracy_score(all_labels, all_preds)
    if epoch % 4 == 0:
        print(f"  {epoch:>4}  {total_loss/len(train_ld):>12.4f}  {val_acc:>10.4f}")

    if val_acc > best_acc:
        best_acc, best_wts = val_acc, copy.deepcopy(model.state_dict())
        patience = 0
    else:
        patience += 1
        if patience >= 6:
            print(f"  Early stop at epoch {epoch}")
            break

model.load_state_dict(best_wts)
print(f"
Best val accuracy: {best_acc:.4f}")

# ── Export for production ─────────────────────────────────────────────
# Save
torch.save({'model_state': model.state_dict(),
             'categories': CATEGORIES,
             'backbone': 'resnet50'}, 'meesho_classifier.pt')

# ONNX export for deployment
dummy = torch.randn(1, 3, 224, 224)
torch.onnx.export(
    model, dummy, 'meesho_classifier.onnx',
    input_names=['image'], output_names=['logits'],
    dynamic_axes={'image': {0: 'batch'}, 'logits': {0: 'batch'}},
    opset_version=17,
)
print("Exported: meesho_classifier.pt + meesho_classifier.onnx")

A health-tech team building an abnormality classifier for chest X-rays reused the exact recipe that had worked well for a product-photo classifier: ResNet50 pretrained on ImageNet, backbone entirely frozen, fine-tune only a new head, 4,000 labelled X-rays. It shipped at 61 percent accuracy on the held-out test set — barely above the 58 percent majority-class baseline.

The cause was domain shift, not a bug. X-rays are single-channel grayscale images forced into 3 identical channels to match ResNet's input shape, their intensity statistics look nothing like natural photographs, and ImageNet's later layers are tuned to recognise things like fur, fabric weave, and wheel spokes — visual cues that simply do not exist in a lung radiograph. The frozen backbone's later layers were producing near-meaningless features for this input distribution, and no amount of head training could recover information that was never extracted in the first place.

The fix: unfreeze layer2 through layer4 with a small learning rate instead of just the head, so the network could re-learn its higher-level features on real X-ray statistics. Accuracy rose to 84 percent. The team later swapped the backbone entirely for one pretrained on radiology images instead of ImageNet, and gained a few more points — evidence that for domains this far from natural photographs, the pretraining dataset matters as much as the fine-tuning strategy.

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

# ── Diagnose domain shift BEFORE a full training run ───────────────────
# Compare frozen-backbone feature statistics on your data vs natural
# photos. Catches "these features won't transfer" early, before you
# burn GPU hours training a head that has nothing useful to work with.

backbone = models.resnet50(weights='IMAGENET1K_V2')
backbone.fc = nn.Identity()   # expose the pooled 2048-d feature vector
backbone.eval()

def feature_stats(images: torch.Tensor):
    """images: (N, 3, 224, 224), already ImageNet-normalised."""
    with torch.no_grad():
        feats = backbone(images)   # (N, 2048)
    dead_frac = (feats == 0).float().mean().item()
    return feats.mean().item(), feats.std().item(), dead_frac

torch.manual_seed(42)

# Reference: activations on natural-photo-like inputs
natural = torch.randn(32, 3, 224, 224) * 0.3 + 0.1
mean_n, std_n, dead_n = feature_stats(natural)

# Your data: chest X-rays, grayscale forced into 3 channels
xray = torch.randn(32, 3, 224, 224) * 0.05 + 0.02
mean_x, std_x, dead_x = feature_stats(xray)

print(f"{'Input distribution':<22} {'feat mean':>10} {'feat std':>10} {'dead units':>12}")
print("─" * 58)
print(f"  {'Natural photos':<20} {mean_n:>10.4f} {std_n:>10.4f} {dead_n*100:>10.1f}%")
print(f"  {'Chest X-rays':<20} {mean_x:>10.4f} {std_x:>10.4f} {dead_x*100:>10.1f}%")

print("
Rule of thumb: a much higher dead-unit fraction on your data than")
print("on natural photos means a large part of the pretrained backbone is")
print("not firing at all on your domain — unfreeze more layers than you")
print("would for an in-domain task, not just the classification head.")

Transfer learning wins when your dataset is small-to-moderate relative to the task's complexity, or when your domain is reasonably close to the pretraining domain — that's most production vision problems, which is why it's the default recipe. But given a truly large, well-labelled dataset (hundreds of thousands of images) in a domain far from natural photographs, training from scratch — or fine-tuning so aggressively it's close to starting over — can match or beat a lightly fine-tuned pretrained model, because the model can learn representations specifically shaped by your data instead of inheriting ImageNet's biases. The right call depends on both dataset size and domain distance, not a blanket rule.

Freezing everything except the head (pure feature extraction) protects you from catastrophic forgetting, but it also caps how well the model can adapt — if your task needs task-specific high-level features the frozen backbone never learned (because ImageNet never needed them), no amount of head training will produce them. "Safer" here just means "less likely to break," not "more likely to perform best." For a dataset of meaningful size with any domain gap from ImageNet, selectively unfreezing later layers with a small learning rate consistently outperforms full freezing — the risk of forgetting is real but manageable, and the risk of under-adapting is often larger.

"Reasonably small" is relative to what the layer already knows, and early backbone layers know a lot — a learning rate that seems conservative for a randomly initialised head (1e-4, say) can still meaningfully disturb pretrained low-level filters over enough epochs, especially without gradient clipping or warmup. Catastrophic forgetting isn't a single threshold you cross; it's a matter of degree that compounds across training steps and depends on how many layers are unfrozen, for how long, and at what relative rate to the head. That's why the standard fix is differential learning rates *and* often gradual unfreezing, not just picking one small number and hoping it's small enough.

Early layers learn genuinely universal low-level patterns — edges, corners, colour gradients — that transfer well almost everywhere. But "transfers well" degrades the further the target domain sits from natural RGB photographs of everyday objects. Medical imaging (X-rays, MRI slices), satellite imagery, and microscopy all have fundamentally different structure, colour statistics, and scale than ImageNet's dog photos and street scenes — a backbone's later layers, tuned specifically to recognise ImageNet-style objects, offer far less of a head start there. In practice this means larger domain shift calls for unfreezing more layers (or fine-tuning more aggressively) than a task that closely resembles natural photography, and in extreme domain shift cases, transfer learning's advantage can shrink to not much more than a decent weight initialisation.

ImageNet top-1 accuracy measures how well a backbone performs at ImageNet's own task and scale — it does not directly predict how well its features transfer to your smaller, different dataset. A large model like ViT-B/16 (86M parameters) generally needs more fine-tuning data to avoid overfitting than a compact ResNet18 (11M parameters), so on a genuinely small dataset the "worse" ImageNet model can win after fine-tuning simply because it doesn't overfit as fast. Backbone choice has to weigh accuracy against your actual dataset size, latency budget, and deployment constraints — "highest ImageNet accuracy" is one input to that decision, not the answer to it.

I'd default to transfer learning whenever labelled data is limited relative to the task (roughly under tens of thousands of images) and the domain isn't wildly different from natural photographs — which covers most production vision problems, since it gets you most of ImageNet's learned representations essentially for free and trains in a fraction of the time. I'd lean toward training from scratch, or very aggressive fine-tuning that's effectively close to it, when I have a large, high-quality dataset in a domain far from ImageNet — satellite imagery or microscopy, for example — where a from-scratch model can learn representations actually shaped by the data instead of inheriting biases that don't apply.

Early convolutional layers learn low-level, local patterns — edges, corners, colour gradients, simple textures — which are useful for recognising essentially any visual input, whether it's a shoe, an X-ray, or a satellite photo. Later layers combine those low-level patterns into progressively more abstract, task-specific representations — by the final layers, filters are effectively tuned to detect things like "dog snout" or "car wheel," which are specific to what the network was trained to classify. That's why the standard fine-tuning recipe freezes or barely touches early layers (their features are already close to optimal for any vision task) while unfreezing later layers with a higher learning rate so they can specialise toward the new task.

It's when fine-tuning updates backbone weights aggressively enough, early enough, that the useful pretrained representations get overwritten before the new head has learned anything useful to build on — you end up with a model that's worse than either the original pretrained backbone or a properly fine-tuned one. It typically comes from using too high a learning rate on backbone layers, especially in the first few epochs while the randomly initialised head is still producing large, noisy gradients that flow back into the backbone. I prevent it with differential learning rates (much smaller lr on backbone layers than the head), often combined with freezing the backbone entirely for the first few epochs so the head stabilises first, and gradient clipping as a safety net.

My first suspicion is domain shift: ImageNet is natural RGB photographs of everyday objects and scenes, and medical images (X-rays, MRI, histopathology slides) differ in colour statistics, texture, and scale in ways that make the backbone's later, more task-specific layers much less useful as a starting point than they'd be for, say, product photos. I'd check whether normalisation statistics match what the backbone expects, then try unfreezing more layers (or a larger portion of the network) than I would for an in-domain task, since less of the pretrained representation transfers cleanly here. I'd also sanity-check the data pipeline itself — grayscale medical images forced into 3-channel RGB, or intensity ranges very different from natural images, are common silent bugs that look exactly like a domain-shift problem.

I'd start with feature extraction as a fast baseline any time the dataset is small (under a few hundred images) or very similar to ImageNet — freeze the entire backbone, train only the new head, and see where accuracy lands. If I have a moderate-to-large dataset or the domain differs meaningfully from ImageNet, I'd move to fine-tuning: unfreeze the later backbone layers (and the head, obviously), keep the earliest layers frozen since their features are already near-universal, and assign a much smaller learning rate to unfrozen backbone layers than to the head — a common pattern is roughly 0.01× the head's learning rate for early unfrozen layers, scaling up toward the head. I'd validate the setup by watching per-layer-group loss behaviour early in training, and consider gradual unfreezing instead if the dataset is on the smaller side.