Replacing classifier head in PyTorch - Practice Problems & Coding Challenges

Challenge - 5 Problems

❓ Predict Output

intermediate

Output of replacing classifier head in a PyTorch model

What is the output shape of the model's final layer after replacing the classifier head with a new linear layer of 10 output features?

PyTorch

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

model = models.resnet18()
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 10)

input_tensor = torch.randn(4, 3, 224, 224)
output = model(input_tensor)
output_shape = output.shape
print(output_shape)

A torch.Size([4, 10])

B torch.Size([4, 1000])

C torch.Size([1, 10])

D torch.Size([4, 512])

💡 Hint

Remember the batch size and the number of output classes in the new classifier head.

Replacing the classifier head changes the output layer to have 10 output features. The batch size is 4, so the output shape is (4, 10).

❓ Model Choice

intermediate

Choosing the correct way to replace classifier head in PyTorch

Which option correctly replaces the classifier head of a pretrained VGG16 model to output 5 classes?

A model.head = nn.Linear(4096, 5)

B model.fc = nn.Linear(512, 5)

C model.classifier = nn.Linear(4096, 5)

D model.classifier[6] = nn.Linear(4096, 5)

💡 Hint

Check the attribute name and index of the classifier layer in VGG16.

VGG16's classifier is a Sequential module, and the last layer is at index 6. Replacing model.classifier[6] correctly changes the output layer.

❓ Hyperparameter

advanced

Effect of freezing layers when replacing classifier head

If you replace the classifier head of a pretrained ResNet50 and freeze all layers except the new head, which statement is true about training?

A All model parameters will update during training.

B Only the new classifier head's parameters will update during training.

C No parameters will update because the model is frozen.

D Only the first convolutional layer will update during training.

💡 Hint

Freezing layers means setting requires_grad to False for those parameters.

Freezing all layers except the new head means only the new head's parameters have requires_grad=True and will update during training.

🔧 Debug

advanced

Debugging error after replacing classifier head

After replacing the classifier head of a pretrained ResNet18 with nn.Linear(512, 20), the model raises a runtime error during training: "size mismatch, m1: [4 x 512], m2: [1000 x 20]". What is the cause?

A The old classifier layer was not replaced properly; the model still uses the original 1000 output features.

B The input tensor batch size is incorrect.

C The new classifier layer has wrong input features; it should be 1000 instead of 512.

D The loss function expects 1000 classes instead of 20.

💡 Hint

Check if the model's classifier attribute was correctly assigned.

The error shows a mismatch between expected input features and the classifier weight matrix, indicating the old classifier was not replaced properly.

🧠 Conceptual

expert

Why replace classifier head instead of retraining entire model?

Why is it common practice to replace only the classifier head of a pretrained model when adapting it to a new task?

A Because retraining the entire model is impossible with pretrained weights.

B Because the classifier head contains all convolutional filters needed for feature extraction.

C Because pretrained layers have learned useful features and retraining only the head saves time and data.

D Because replacing the head increases the model size and improves accuracy automatically.

💡 Hint

Think about transfer learning and feature reuse.

Pretrained layers capture general features; retraining only the head adapts to new classes efficiently without needing large data or long training.

Practice

1. What is the main reason to replace the classifier head in a pretrained PyTorch model?

easy

A. To adapt the model to a new task with different output classes

B. To speed up the training by removing layers

C. To reduce the model size by deleting layers

D. To change the input image size the model accepts

Solution

1. Step 1: Understand the classifier head role

   The classifier head is the last layer that decides the output classes based on learned features.

2. Step 2: Reason about adapting to new tasks

   Replacing the classifier head allows the model to output predictions for new classes different from the original training.

3. Final Answer:

   To adapt the model to a new task with different output classes -> Option A

4. Quick Check:

   Classifier head replacement = new task adaptation [OK]

Hint: Classifier head controls output classes, replace for new tasks [OK]

Common Mistakes:

• Thinking replacing head changes input size
• Assuming it reduces model size significantly
• Believing it speeds up training by removing layers

2. Which of the following is the correct way to replace the classifier head of a pretrained ResNet model in PyTorch for 10 output classes?

easy

A. model.fc = nn.Linear(2048, 10)

B. model.classifier = nn.Linear(2048, 10)

C. model.fc = nn.Linear(512, 10)

D. model.head = nn.Linear(512, 10)

Solution

1. Step 1: Identify ResNet classifier attribute

   ResNet models use model.fc as the classifier head.

2. Step 2: Check input feature size for ResNet

   ResNet50 and similar have 2048 features before the classifier, so input size is 2048.

3. Final Answer:

   model.fc = nn.Linear(2048, 10) -> Option A

4. Quick Check:

   ResNet classifier = model.fc with 2048 input features [OK]

Hint: ResNet classifier is model.fc with 2048 input features [OK]

Common Mistakes:

• Using wrong attribute like model.classifier or model.head
• Using wrong input size like 512 instead of 2048
• Confusing ResNet with other models like VGG

3. Given the code below, what will be the output shape of the model's final layer after replacement?

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

model = models.resnet18(pretrained=True)
model.fc = nn.Linear(512, 5)

input_tensor = torch.randn(1, 3, 224, 224)
output = model(input_tensor)
print(output.shape)

medium

A. torch.Size([1, 1000])

B. torch.Size([1, 512])

C. torch.Size([1, 5])

D. torch.Size([3, 224, 224])

Solution

1. Step 1: Understand the replaced classifier output size

   The new classifier layer outputs 5 values per input (5 classes).

2. Step 2: Check input batch size and output shape

   Input batch size is 1, so output shape is (1, 5).

3. Final Answer:

   torch.Size([1, 5]) -> Option C

4. Quick Check:

   Output shape = (batch_size, output_classes) = (1, 5) [OK]

Hint: Output shape matches batch size and new class count [OK]

Common Mistakes:

• Expecting original 1000 classes output
• Confusing feature size with output size
• Misreading input tensor shape as output

4. You tried replacing the classifier head of a pretrained model with model.fc = nn.Linear(1024, 10) but got a runtime error during training. What is the most likely cause?

medium

A. The model.fc attribute does not exist in pretrained models

B. The output size 10 is too large for the model

C. You forgot to call model.eval() before training

D. The input feature size 1024 does not match the model's actual output features

Solution

1. Step 1: Check input feature size for classifier

   The input size to the new Linear layer must match the output features of the previous layer.

2. Step 2: Identify mismatch causing runtime error

   If 1024 is incorrect, the model will raise size mismatch errors during forward pass.

3. Final Answer:

   The input feature size 1024 does not match the model's actual output features -> Option D

4. Quick Check:

   Input size mismatch causes runtime error [OK]

Hint: Match Linear input size to previous layer output features [OK]

Common Mistakes:

• Assuming output size causes error
• Confusing eval mode with training errors
• Thinking model.fc is missing in pretrained models

5. You want to fine-tune a pretrained ResNet50 on a dataset with 15 classes. Which code snippet correctly replaces the classifier head and freezes all layers except the new head?

hard

A. model = models.resnet50(pretrained=True)
model.fc = nn.Linear(2048, 15)
for param in model.parameters():
param.requires_grad = False

B. model = models.resnet50(pretrained=True)
for param in model.parameters():
param.requires_grad = False
model.fc = nn.Linear(2048, 15)

C. model = models.resnet50(pretrained=True)
for param in model.fc.parameters():
param.requires_grad = False
model.fc = nn.Linear(2048, 15)

D. model = models.resnet50(pretrained=True)
model.fc = nn.Linear(512, 15)
for param in model.parameters():
param.requires_grad = True

Solution

1. Step 1: Freeze all existing model parameters

   Set param.requires_grad = False for all parameters to prevent updates during training.

2. Step 2: Replace classifier head with correct input/output sizes

   ResNet50's classifier input size is 2048; output size is 15 for new classes.

3. Step 3: Ensure new head parameters are trainable

   By replacing model.fc after freezing, new layer parameters default to requires_grad=True.

4. Final Answer:

   Freeze all params, then replace head with nn.Linear(2048, 15) -> Option B

5. Quick Check:

   Freeze old layers, replace head with correct sizes [OK]

Hint: Freeze before replacing head to keep new layer trainable [OK]

Common Mistakes:

• Freezing after replacing head disables new layer training
• Using wrong input size 512 instead of 2048
• Not freezing any layers when fine-tuning