Transfer Learning: Fine-Tuning Pre-Trained Vision Models

Learn Transfer Learning. Understand how to reuse features from pre-trained vision models (ResNet, VGG) and fine-tune them in PyTorch.

Introduction

Imagine you want to learn how to drive a truck. You do not start by learning how a combustion engine works, nor do you spend years learning how to balance on a bicycle again. You take your existing knowledge of driving a car (steering, braking, traffic laws) and adapt it to the truck.

In deep learning, this is called Transfer Learning. Instead of training a massive neural network from scratch—which requires millions of images and weeks of GPU compute—we take a model pre-trained on a massive dataset (like ImageNet) and adapt its learned features to solve a new, specific task.

What You Will Learn

  • The core concept and philosophy of Transfer Learning.
  • The difference between Feature Extraction and Fine-Tuning.
  • What it means to Freeze Layers.
  • How to implement Transfer Learning in PyTorch.

Why This Topic Matters

Almost no one in the industry trains deep vision models from scratch. It is too expensive and data-intensive. Transfer learning allows you to build a world-class image classifier (e.g., classifying rare skin diseases) using as few as 100 training images, completing training in just a few minutes on a standard CPU.

Prerequisites

Detailed Explanation

A deep CNN trained on ImageNet (1.2 million images, 1,000 classes) learns hierarchical features:

  • Early layers learn general features (edges, corners).
  • Middle layers learn textures and shapes.
  • Final layers learn class-specific parts (dog ears, car wheels).

Since edges, shapes, and textures are universal to all images, we can reuse the early and middle layers for our new dataset (e.g., medical scans). We only need to replace and train the final classification layer.

Two Strategies of Transfer Learning

graph TD
    A[Transfer Learning Strategies]
    A --> B[1. Feature Extraction <br> Freeze all pre-trained layers <br> Only train the new output classifier]
    A --> C[2. Fine-Tuning <br> Unfreeze some/all pre-trained layers <br> Train the whole network with a tiny learning rate]

1. Feature Extraction (Recommended for small datasets)

We treat the pre-trained model as a fixed feature extractor.

  • We freeze all the convolutional layers (meaning their weights do not change during backpropagation).
  • We remove the original final classifier (which has 1,000 outputs) and replace it with a new fully connected layer matching our target class count (e.g., 2 outputs for "Cancer" vs. "Normal").
  • We only train the weights of this new final layer.

2. Fine-Tuning (Recommended for medium/large datasets)

We don't just replace the final layer; we also unfreeze some (or all) of the convolutional layers.

  • We train the entire network using a very small learning rate (e.g., $10^{-5}$).
  • This allows the model to "fine-tune" its pre-trained filters to match the specific details of our new dataset, without destroying the general features it already knows.

When to Use Which Strategy?

| Dataset Size | Similarity to Pre-trained Data | Recommended Strategy | | :--- | :--- | :--- | | Very Small | High similarity (e.g., cats vs. dogs) | Feature Extraction (prevent overfitting). | | Very Small | Low similarity (e.g., satellite scans) | Difficult. Try extracting from early layers only. | | Large | High similarity | Fine-Tuning (model will customize and peak in accuracy). | | Large | Low similarity | Fine-Tuning or training from scratch. |

Visual Diagram (Mermaid)

graph LR
    subgraph Pre-Trained Network
    A[Conv Layers: FROZEN <br> Keeps learned weights] --> B[Dense Classifier: REPLACED]
    end
    B -->|Only train this| C[New Class Labels]
    style A fill:#3B82F6,stroke:#fff,color:#fff
    style C fill:#10B981,stroke:#fff,color:#fff

Python Code Examples

We will implement a transfer learning workflow in PyTorch using a pre-trained ResNet-18 model.

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

# 1. Load a pre-trained ResNet-18 model
# weights=ResNet18_Weights.DEFAULT downloads the ImageNet trained weights
model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)

# 2. Freeze all convolutional layers (disable gradients calculation)
for param in model.parameters():
    param.requires_grad = False

# 3. Inspect the original final fully connected (fc) layer
print("Original FC Layer structure:")
print(model.fc) 
# Expected: Linear(in_features=512, out_features=1000, bias=True)

# 4. Replace final layer with a new one for 2 target classes (e.g., Cat vs Dog)
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 2) # This new layer has requires_grad=True by default

print("\nUpdated FC Layer structure:")
print(model.fc)

# 5. Define optimizer
# We pass ONLY the parameters of the new final layer to the optimizer
optimizer = torch.optim.SGD(model.fc.parameters(), lr=0.001, momentum=0.9)

Industry Use Cases

  • Autonomous Drones: Fine-tuning pre-trained ResNet backbones to detect cracks on high-voltage power lines using drone footage.
  • Medical Diagnostics: Reusing models trained on millions of consumer photos to classify lung infections in chest X-rays.
  • E-commerce Product Search: Taking a general feature extractor to map images of shoes into vector spaces for catalog recommendation.

Summary

Transfer learning accelerates deep learning development by reusing feature maps trained on massive general datasets. By freezing spatial layers and swapping final fully connected classifiers, engineers can train high-performance classifiers on small datasets in minutes, bypassing the high computing cost of training from scratch.

Next Topic

Congratulations! You have completed the Computer Vision module. Now, how do we deploy these deep models to production servers so they can serve real-time API requests? Let's check: Introduction to MLOps.

← Back to Syllabus