Implementing Transfer Learning with Pre-trained Models in PyTorch: A Complete Guide

To implement transfer learning in PyTorch, load a pre-trained model from torchvision, freeze its convolutional layers, replace the final classifier with a new output layer matching your task, then train only the new head on your dataset.

The microsoft/AI-For-Beginners repository provides a hands-on demonstration of this workflow in TransferLearningPyTorch.ipynb. This guide walks through extracting the core patterns so you can apply transfer learning to your own image classification tasks with minimal data and computational overhead.

---

Why Transfer Learning Matters for Computer Vision

Training deep convolutional networks from scratch demands massive datasets and GPU resources. Transfer learning solves this by reusing feature extractors pre-trained on ImageNet (1.2 million images, 1000 classes). The lower layers learn universal patterns—edges, textures, shapes—while only the final classifier needs adaptation to your specific classes.

The repository demonstrates this with a classic cats-vs-dogs binary classification task, achieving strong results by training just a single linear layer atop a frozen VGG-16 backbone.

---

Loading and Inspecting Pre-trained Models

PyTorch's torchvision provides ready-to-use architectures with ImageNet weights. The notebook starts by loading VGG-16 and examining its structure at line 90:

import torch
import torchvision

device = 'cuda' if torch.cuda.is_available() else 'cpu'

vgg = torchvision.models.vgg16(pretrained=True).to(device)

Key architectural components to understand:

• vgg.features — The convolutional stack (conv → ReLU → pooling) that outputs 512×7×7 feature maps
• vgg.classifier — The sequential fully-connected layers that map to 1000 ImageNet classes

Inspecting the model with torchinfo.summary() reveals the parameter distribution: ~138 million total parameters, with the feature extractor comprising the bulk.

---

Freezing the Feature Extractor

To preserve learned representations, freeze convolutional parameters by setting requires_grad=False. This is implemented at line 140:

for param in vgg.features.parameters():
    param.requires_grad = False

This prevents gradients from flowing backward through the feature layers during backpropagation, effectively locking the visual vocabulary learned from ImageNet. Only parameters with requires_grad=True will update during training.

---

Replacing the Classifier Head

The original 1000-way ImageNet classifier must be swapped for a task-appropriate output layer. For cats-vs-dogs (2 classes), the notebook replaces the entire classifier with a single linear layer at line 138:

vgg.classifier = torch.nn.Linear(25088, 2).to(device)

The input dimension 25088 matches the flattened output of VGG-16's final pooling layer (512 channels × 7 × 7 spatial dimensions). For different architectures, always verify this value via torchinfo.summary() or by inspecting model.classifier[0].in_features.

---

Training the New Classifier

With the feature extractor frozen and classifier replaced, training focuses exclusively on the new 2-way head. The notebook leverages helper utilities from [pytorchcv.py](https://github.com/microsoft/AI-For-Beginners/blob/main/lessons/4-ComputerVision/08-TransferLearning/pytorchcv.py) (imported at line 29):

from pytorchcv import train

history = train(vgg, train_loader, test_loader)

This phase trains only the ~50,000 parameters in the new linear layer—orders of magnitude fewer than full network training—enabling fast convergence even on modest hardware.

---

Fine-Tuning the Entire Network

When the target domain differs substantially from ImageNet, fine-tuning adapts the feature extractor to new visual characteristics. The notebook demonstrates this two-stage approach:

Step 1: Train the classifier (feature extraction)

Complete as shown above to establish a reasonable starting point.

Step 2: Unfreeze and train with reduced learning rate

At line 155, the convolutional layers are unfrozen:

from pytorchcv import train_long

# Unfreeze feature layers

for param in vgg.features.parameters():
    param.requires_grad = True

# Train entire network with small learning rate

train_long(vgg, train_loader, test_loader,
           loss_fn=torch.nn.CrossEntropyLoss(),
           epochs=1,
           lr=0.0001)  # 10-100× smaller than initial training

Critical considerations for fine-tuning:

• Use a significantly lower learning rate (typically 0.0001 vs. 0.001) to avoid overwriting pre-trained features
• Consider differential learning rates: even lower rates for early layers, moderate rates for deeper layers, standard rates for the classifier
• Unfreeze in stages rather than all at once for very large models

---

Working with Alternative Architectures: ResNet-18

The same pattern applies across architectures. The notebook demonstrates ResNet-18 loading at line 207:

resnet = torchvision.models.resnet18(pretrained=True).to(device)

# Freeze all parameters

for param in resnet.parameters():
    param.requires_grad = False

# Replace final fully-connected layer

resnet.fc = torch.nn.Linear(resnet.fc.in_features, 2).to(device)

ResNet organizes layers differently than VGG:

• Feature extractor: conv1 through layer4 (residual blocks with skip connections)
• Classifier: single fc layer rather than a sequential stack

Always verify your target architecture's component names via print(model) or documentation before freezing.

---

Dataset Preparation and Preprocessing

Pre-trained models expect specific input normalization. The notebook applies ImageNet statistics at transformation time:

std_normalize = transforms.Normalize(
    mean=[0.485, 0.456, 0.406],
    std=[0.229, 0.224, 0.225]
)

transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),  # matches ImageNet training size

    transforms.ToTensor(),
    std_normalize,
])

ImageNet normalization is essential—the pre-trained weights assume this input distribution. Deviating without retraining the feature extractor produces degraded performance.

---

GPU Acceleration and Device Handling

The notebook consistently handles device placement, as shown at line 124:

device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = model.to(device)

Key patterns:

• Move models to device immediately after loading
• Ensure input tensors are on the same device during training/inference
• Use torch.no_grad() context for inference to save memory

---

Summary

• Load pre-trained models via torchvision.models with pretrained=True to access ImageNet weights
• Freeze feature extractors by setting requires_grad=False on convolutional parameters to preserve learned visual features
• Replace classifier heads with task-appropriate output dimensions, matching the flattened feature dimension of your chosen architecture
• Train in stages: classifier-only first for fast convergence, then optional fine-tuning with reduced learning rates
• Normalize inputs using ImageNet statistics (mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) to match pre-training conditions
• Reuse patterns across architectures—VGG, ResNet, and others share the same fundamental structure of frozen backbone + trainable head

---

Frequently Asked Questions

How do I determine the input size for my replacement classifier layer?

Inspect the output shape of the feature extractor using torchinfo.summary() or by passing a dummy tensor through model.features. For VGG-16, this yields 512×7×7=25088; for ResNet-18, it's 512 (global average pooling collapses spatial dimensions). Always verify rather than hardcoding, as different model variants vary.

What's the difference between feature extraction and fine-tuning?

Feature extraction freezes all pre-trained layers and trains only the new classifier—fast, low compute, suitable when target data is limited or similar to ImageNet. Fine-tuning unfreezes earlier layers and trains the full network with reduced learning rates—better when target domain differs significantly (medical imaging, satellite imagery) and more data is available.

When should I use a smaller learning rate for fine-tuning?

Always use 10-100× lower learning rates when fine-tuning pre-trained weights (typically 1e-4 vs. 1e-3). Pre-trained features are already optimized; aggressive updates destroy useful representations. The train_long helper in the repository accepts an explicit lr parameter for this control.

Can I apply this same pattern to other modalities like NLP or audio?

Yes—the transfer learning principle of freeze base → replace head → train → optionally fine-tune applies across domains. The repository demonstrates this pattern in [examples/04-text-sentiment.py](https://github.com/microsoft/AI-For-Beginners/blob/main/examples/04-text-sentiment.py) for text classification using pre-trained embeddings, and similar approaches work for audio spectrograms with pre-trained speech models.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}