Transfer Learning for Vision

Build the mental model

The core insight behind transfer learning is that the early and middle layers of a CNN trained on a large dataset (like ImageNet) end up producing general-purpose visual features — edges, textures, shapes, color gradients — that aren't specific to the original task. It's an empirical fact that these general features transfer across nearly all natural-photo vision tasks: an edge detector learned for classifying cats turns out to be just as useful for classifying cars. Freezing parameters with requires_grad = False means the backward pass skips computing gradients for those layers entirely, which makes training both faster and less memory-hungry, and it also sidesteps the overfitting that would happen if you tried to train an entire deep network from scratch on a small dataset.

Replacing model.fc with a fresh Linear layer only swaps out the network's 'decision head' — the convolutional backbone stays exactly as it was trained on ImageNet's 1000 classes, and only the classifier head gets replaced to match the target task's number of classes. Since the only trainable parameters left are the weights and biases of that new layer, you no longer need an ImageNet-scale dataset — a comparatively small labeled image set is enough to reach good accuracy.

Connect it to a real scenario

Suppose the Tutorial Platform needs a moderation feature that flags inappropriate content in user-uploaded profile avatars. Instead of needing thousands of labeled avatar images to train a classifier from scratch, freezing a pretrained ResNet18 and swapping only its fc layer for a binary (appropriate/inappropriate) output lets you get a working prototype with a much smaller labeled avatar dataset.

The early/middle layers of a pretrained CNN produce general visual features that transfer well across many vision tasks.
Freezing with requires_grad = False speeds up training and reduces overfitting risk on a small target dataset.
Replacing model.fc with a new Linear layer only swaps the decision head to match the target task's number of classes.

Try the working example

python

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

model = models.resnet18(weights=ResNet18_Weights.DEFAULT)

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

# Replace the final classification layer for a new task with 5 classes
num_classes = 5
model.fc = nn.Linear(model.fc.in_features, num_classes)

trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
frozen = sum(p.numel() for p in model.parameters() if not p.requires_grad)

print(f"trainable parameters: {trainable}")
print(f"frozen parameters: {frozen}")
print(f"total parameters: {trainable + frozen}")

You should see

Prints three lines: 'trainable parameters:', 'frozen parameters:', and 'total parameters:'. The trainable count is exactly the size of the newly-created fc Linear layer (512 in_features × 5 classes + 5 bias terms = 2,565 parameters); the frozen count is everything else in the ResNet18 backbone (roughly 11 million-plus parameters) — so trainable parameters make up a tiny fraction of the total.

5-minute try-it

Try 'partial fine-tuning' by also setting requires_grad = True on the parameters inside model.layer4 (ResNet18's final convolutional block), in addition to the new fc layer, then print how much the trainable parameter count increases.

One important caution

Downloading pretrained weights via ResNet18_Weights.DEFAULT requires internet access the first time; running this code offline (without a cached weights file) will error out.

Passing model.parameters() directly to the optimizer after replacing fc — instead of filtering with something like filter(lambda p: p.requires_grad, model.parameters()) — includes the frozen parameters in the optimizer's state and wastes memory even though they never update.

PyTorch Tutorials — Transfer Learning for Computer Vision — Computer Vision