Replacing classifier head in PyTorch - Model Pipeline Trace

Model Pipeline - Replacing classifier head

This pipeline shows how to replace the last part of a pre-trained model (called the classifier head) with a new one. This helps the model learn to classify new categories while keeping the earlier learned features.

Data Flow - 4 Stages

1Load pre-trained model

N/A→Load a model trained on a large dataset (e.g., ImageNet)→Model with original classifier head for 1000 classes

ResNet50 with final layer output size 1000

↓

2Remove original classifier head

Model with final layer output size 1000→Remove or ignore the last fully connected layer→Model feature extractor output size 2048

ResNet50 without final fc layer outputs 2048 features

↓

3Add new classifier head

2048 features→Add a new fully connected layer for new number of classes→Model with new classifier head output size 10

New fc layer outputs 10 classes for new task

↓

4Train model on new data

Images with labels for 10 classes→Train only the new classifier head or fine-tune whole model→Trained model adapted to new classes

Model learns to classify 10 new categories

Training Trace - Epoch by Epoch

Loss
2.0 | *       
1.5 |  *      
1.0 |   *     
0.5 |     *   
0.0 |       * 
     --------
     Epochs

Epoch	Loss ↓	Accuracy ↑	Observation
1	1.8	0.40	Loss starts high, accuracy low as model begins learning new classes
2	1.2	0.60	Loss decreases, accuracy improves as classifier head adapts
3	0.8	0.75	Model shows good learning progress
4	0.6	0.82	Loss continues to decrease, accuracy rises
5	0.5	0.87	Training converges with good accuracy

Prediction Trace - 5 Layers

Layer 1: Input image

Layer 2: Feature extractor (all layers except classifier head)

Layer 3: New classifier head (fully connected layer)

Layer 4: Softmax activation

Layer 5: Prediction

Model Quiz - 3 Questions

Why do we replace the classifier head in a pre-trained model?

A To reduce the input image size

B To adapt the model to new classes different from the original

C To make the model run faster

D To remove all learned features

Key Insight

Replacing the classifier head allows a pre-trained model to learn new classes efficiently by keeping useful features and only training a small new part. This speeds up training and improves performance on new tasks.