Replacing classifier head in PyTorch

Introduction

We replace the classifier head to change what the model predicts. This helps when using a pre-trained model for a new task with different output classes.

You want to use a pre-trained model but with a different number of output classes.

You need to adapt a model trained on one dataset to work on another dataset.

You want to fine-tune only the last layer for faster training.

You want to experiment with different classifier designs on top of a fixed feature extractor.

Syntax

PyTorch

model.classifier = torch.nn.Linear(in_features, out_features)

model.classifier is the part of the model that makes final predictions.

in_features must match the output size of the previous layer.

Examples

Replace classifier with a linear layer that outputs 10 classes, assuming input features are 512.

PyTorch

model.classifier = torch.nn.Linear(512, 10)

For models like ResNet, the classifier is often called fc. This replaces it to output 5 classes.

PyTorch

model.fc = torch.nn.Linear(2048, 5)

Sample Model

This code loads a pre-trained ResNet18 model, replaces its classifier head to output 3 classes, and runs a dummy input through it. It prints the original classifier size, the output shape, and the output values.

PyTorch

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

# Load a pre-trained ResNet18 model
model = models.resnet18(pretrained=True)

# Check original classifier (fc) output features
print(f'Original classifier output features: {model.fc.out_features}')

# Replace the classifier head to output 3 classes
model.fc = nn.Linear(model.fc.in_features, 3)

# Create dummy input tensor (batch size 1, 3 color channels, 224x224 image)
dummy_input = torch.randn(1, 3, 224, 224)

# Get model output
output = model(dummy_input)

# Print output shape and values
print(f'Output shape: {output.shape}')
print(f'Output values: {output}')

Important Notes

Always match the input features of the new classifier to the previous layer's output size.

Replacing the classifier head is common in transfer learning.

After replacing, you usually need to train or fine-tune the new head.

Summary

Replacing the classifier head lets you adapt a model to new tasks.

Use the correct input and output sizes for the new layer.

This is a key step in transfer learning with PyTorch models.