PyTorch Transfer Learning Intro

Transfer learning takes a model already trained on a large dataset and reuses it for a new, different task. Instead of training from random weights, you start from weights that already understand edges, shapes, textures, or language patterns. This gives your model a head start and dramatically reduces the amount of data and compute needed.

The Core Idea

Without transfer learning:
  Task: classify 500 medical X-ray images (cat vs dog-sized dataset)
  Start: random weights
  Problem: 500 images is far too few to learn edge detectors, shapes,
           textures — the model overfits immediately

With transfer learning:
  Take ResNet50 trained on 1.2 million ImageNet images
  It already knows: edges → corners → textures → shapes → objects
  Reuse those early layers for your task
  Train only the final classification head on your 500 images
  Result: good accuracy even with limited data

Two Transfer Learning Strategies

Strategy 1: Feature Extraction
  → Freeze all pretrained layers (set requires_grad=False)
  → Replace and train only the final classification head
  → Use when: dataset is small, task is similar to original

Strategy 2: Fine-Tuning
  → Start from pretrained weights
  → Unfreeze some or all layers
  → Train at a very low learning rate
  → Use when: you have more data, task differs from original

Why Pretrained Models Transfer Well

Deep networks trained on ImageNet learn features in a hierarchy. Early layers detect edges and color gradients — features useful for any image task. Middle layers detect textures and simple shapes — broadly applicable. Late layers detect object-specific features — the part most specific to the original task.

Layer depth     Learned features          Transferability
──────────────────────────────────────────────────────────
First few       Edges, gradients          Universal — keep these
Middle          Textures, shapes          Broadly useful — usually keep
Last few        Task-specific patterns    Replace these for new task

Feature Extraction Workflow

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

# Load a pretrained ResNet18
backbone = models.resnet18(pretrained=True)

# Freeze all parameters — backbone is just a feature extractor
for param in backbone.parameters():
    param.requires_grad = False

# Replace the final classification layer
num_classes = 5    # your new task has 5 classes
in_features = backbone.fc.in_features    # 512 for ResNet18
backbone.fc = nn.Linear(in_features, num_classes)

# Only the new fc layer has requires_grad=True
trainable = sum(p.numel() for p in backbone.parameters() if p.requires_grad)
print(f"Trainable parameters: {trainable}")   # only ~2,565

Fine-Tuning Workflow

backbone = models.resnet18(pretrained=True)

# Replace head
backbone.fc = nn.Linear(backbone.fc.in_features, 5)

# Unfreeze everything — train all layers at a low rate
for param in backbone.parameters():
    param.requires_grad = True

# Use a very small learning rate to preserve pretrained features
optimizer = torch.optim.Adam(backbone.parameters(), lr=1e-4)

When to Use Each Strategy

Dataset size        Task similarity     Strategy
─────────────────────────────────────────────────────
Small (<5k)         Similar             Feature extraction
Small (<5k)         Different           Feature extraction + more dropout
Medium (5k–50k)     Similar             Fine-tune last few layers
Medium              Different           Fine-tune most layers
Large (>50k)        Any                 Full fine-tuning or train from scratch

Popular Pretrained Models in TorchVision

Model          Parameters   Use Case
──────────────────────────────────────────────────
ResNet18       11M          Fast, good baseline
ResNet50       25M          Strong accuracy, common choice
EfficientNetB0 5.3M         Efficient, mobile-friendly
VGG16          138M         Simple architecture, heavy
MobileNetV3    5.4M         Lightweight, fast inference
ViT-B/16       86M          Vision Transformer, state of the art

Transfer Learning in NLP

Transfer learning is even more dominant in NLP than in computer vision. Models like BERT, RoBERTa, and GPT are trained on billions of words and then fine-tuned for specific tasks — sentiment analysis, question answering, named entity recognition — with very little task-specific data.

The Hugging Face transformers library integrates directly with PyTorch and provides hundreds of pretrained language models with a simple fine-tuning interface.

Summary

Transfer learning reuses weights from a model trained on a large dataset to jumpstart training on a new, smaller task. Feature extraction freezes the pretrained backbone and trains only a new classification head — fast and effective for small datasets. Fine-tuning unfreezes some or all layers and trains at a low learning rate — better when you have more data or the new task differs significantly from the original. Early layers of pretrained models learn universal features that transfer to almost any vision task. Start with feature extraction, then try fine-tuning if accuracy is insufficient.