Project: Transfer Learning Classifier

Build the mental model

The core idea of transfer learning is that a pretrained resnet18 has already learned generic visual features — edges, textures, shapes — from a huge number of ImageNet images. Freezing every backbone parameter by setting requires_grad = False stops gradients from being computed for those weights, which both saves compute and protects those learned general features from catastrophic forgetting, the risk of a tiny fine-tuning dataset overwriting what a much larger dataset originally taught the network. Replacing the final fully connected layer (model.fc) with a fresh Linear layer swaps out the head that was specialized for ImageNet's 1000 classes for one sized to your actual problem — here, 4 classes.

Building the optimizer over model.fc.parameters() alone, instead of model.parameters(), sharply shrinks the number of trainable parameters from the backbone's millions down to the new fc layer's few thousand — the core promise of transfer learning, that you can get reasonable performance from very few examples and very few training steps. Checking named_parameters() and confirming .grad is None on every non-fc parameter after training is a genuinely good debugging habit: it verifies freezing actually worked, rather than just trusting that setting requires_grad did what you intended.

Connect it to a real scenario

The Tutorial Platform runs a lesson-screenshot quality classifier that flags blurry or low-quality screenshots instructors upload into lessons. Because labeled examples are scarce — they come from a small batch of manual review — a frozen resnet18 backbone with a fine-tuned fc head is exactly what makes this feasible: transfer learning gets a usable classifier out of a handful of labeled examples instead of requiring a large labeled dataset from scratch.

• Setting requires_grad = False freezes the pretrained backbone, and building the optimizer over only the new layer's parameters sharply cuts the trainable parameter count.
• Replacing model.fc must happen before the freezing loop runs, or the newly added layer gets frozen too, right alongside the backbone.
• Checking .grad is None on backbone parameters after training gives objective confirmation that freezing actually behaved as intended.

Try the working example

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision.models import resnet18, ResNet18_Weights

torch.manual_seed(0)

num_classes = 4
model = resnet18(weights=ResNet18_Weights.DEFAULT)

# Freeze all pretrained layers first
for param in model.parameters():
    param.requires_grad = False

# Then replace the final layer -- its params require_grad=True by default
in_features = model.fc.in_features
model.fc = nn.Linear(in_features, num_classes)

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.fc.parameters(), lr=1e-3)

# Synthetic images shaped like the inputs resnet18 expects
num_samples = 16
images = torch.randn(num_samples, 3, 224, 224)
labels = torch.randint(0, num_classes, (num_samples,))

model.train()
losses = []
for step in range(5):
    optimizer.zero_grad()
    outputs = model(images)
    loss = criterion(outputs, labels)
    loss.backward()
    optimizer.step()
    losses.append(loss.item())
    print(f"Step {step+1}/5 - loss: {loss.item():.4f}")

# Confirm only the new fc layer received gradients
frozen_untouched = all(
    p.grad is None for name, p in model.named_parameters() if not name.startswith("fc")
)
fc_has_grad = all(p.grad is not None for p in model.fc.parameters())
print(f"Backbone parameters untouched (grad is None): {frozen_untouched}")
print(f"New fc layer parameters have gradients: {fc_has_grad}")

(The first run may print a download progress log for the pretrained resnet18 weights, which requires internet access.) Then five lines print, "Step 1/5 - loss: X.XXXX" through "Step 5/5 - loss: X.XXXX". Loss starts near ln(4)≈1.386, the random-guess value for 4 classes, and because the single fc layer only has to fit 16 samples on top of fixed 512-dimensional features (heavily overparameterized for the task), it should drop noticeably within just 5 Adam steps. The final two lines print "Backbone parameters untouched (grad is None): True" and "New fc layer parameters have gradients: True".

PyTorch Tutorials — Transfer Learning for Computer Vision — Computer Vision