Replacing the classifier head

Setup#

pip install torch torchvision

Written and run against torch 2.13.0 (CPU) and torchvision 0.28.0. Both scripts use weights=None, so nothing is downloaded — the layer names and shapes are identical with real weights.

Step one: find the head, do not guess it#

import torch
from torchvision import models

# Every family names its final layer differently. Look, do not guess.
for name in ["resnet18", "efficientnet_b0", "mobilenet_v3_small", "convnext_tiny", "vit_b_16"]:
    m = getattr(models, name)(weights=None)
    last = list(m.named_modules())[-1]
    print(f"{name:<20} last module: {last[0]:<18} {last[1]}")

resnet18             last module: fc                 Linear(in_features=512, out_features=1000, bias=True)
efficientnet_b0      last module: classifier.1       Linear(in_features=1280, out_features=1000, bias=True)
mobilenet_v3_small   last module: classifier.3       Linear(in_features=1024, out_features=1000, bias=True)
convnext_tiny        last module: classifier.2       Linear(in_features=768, out_features=1000, bias=True)
vit_b_16             last module: heads.head         Linear(in_features=768, out_features=1000, bias=True)

Five architectures, five different attribute paths, five different in_features. Any tutorial that hardcodes model.fc and 512 works for exactly one of these five.

Notice classifier.1 and classifier.3. In those models classifier is a Sequential holding dropout, normalisation or reshaping layers ahead of the final Linear:

efficientnet_b0.classifier = Sequential(Dropout(p=0.2), Linear(1280, 1000))
convnext_tiny.classifier   = Sequential(LayerNorm2d(768), Flatten(), Linear(768, 1000))

Assigning model.classifier = nn.Linear(...) replaces the whole block. On EfficientNet that deletes the dropout. On ConvNeXt it deletes a LayerNorm and a Flatten, and the forward pass fails.

Step two: one function that handles all of them#

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

def replace_head(model, num_classes):
    """Swap the last nn.Linear for a fresh one with num_classes outputs."""
    name, old = [(n, m) for n, m in model.named_modules() if isinstance(m, nn.Linear)][-1]
    parent_name, _, attr = name.rpartition(".")
    parent = model.get_submodule(parent_name) if parent_name else model
    setattr(parent, attr, nn.Linear(old.in_features, num_classes))
    return name, old.in_features

torch.manual_seed(0)
x = torch.randn(2, 3, 224, 224)
for name in ["resnet18", "efficientnet_b0", "convnext_tiny", "vit_b_16"]:
    m = getattr(models, name)(weights=None)
    where, feats = replace_head(m, num_classes=5)
    m.eval()
    with torch.no_grad():
        out = m(x)
    print(f"{name:<18} replaced {where:<14} ({feats} features in) -> output {tuple(out.shape)}")

resnet18           replaced fc             (512 features in) -> output (2, 5)
efficientnet_b0    replaced classifier.1   (1280 features in) -> output (2, 5)
convnext_tiny      replaced classifier.2   (768 features in) -> output (2, 5)
vit_b_16           replaced heads.head     (768 features in) -> output (2, 5)

Line by line, for the parts that are not obvious at a glance#

name.rpartition(".") splits "classifier.1" into parent path "classifier" and attribute "1". For "fc" the parent path comes back empty, which is why the next line falls back to model itself.

model.get_submodule(parent_name) walks a dotted path and returns the module. It handles Sequential indices as attribute names, which is why setattr(parent, "1", ...) works on a Sequential.

setattr(parent, attr, nn.Linear(...)) is a real assignment on an nn.Module. PyTorch registers the new layer, drops the old, and updates model.parameters(). See module registration and parameters for why plain attribute assignment is enough.

isinstance(m, nn.Linear) and taking [-1] relies on named_modules() returning definition order. That holds for every torchvision classification model. It would not hold for a model whose head is a convolution, which is normal in segmentation.

Common mistakes#

Hardcoding the feature width. nn.Linear(512, 4) on an EfficientNet-B0 raises a shape error at the first forward pass, because 1280 numbers arrive. Fix: read in_features from the layer you are replacing.

Replacing a Sequential classifier wholesale. model.classifier = nn.Linear(1280, 4) on EfficientNet compiles and trains, and quietly removes the dropout that was regularising it. Fix: replace model.classifier[-1] alone.

Labels outside the new class range. Build a 4-class head, feed a label of 4, and cross-entropy raises:

IndexError : Target 4 is out of bounds.

Fix: map your label strings to 0 .. num_classes-1 once, keep the mapping in a file, and assert labels.max() < num_classes before the first epoch.

Putting a softmax inside the head. nn.Sequential(nn.Linear(512, 4), nn.Softmax(dim=1)) then fed to nn.CrossEntropyLoss applies softmax twice. Training still moves, badly. Fix: heads output raw scores. CrossEntropyLoss expects raw scores and does the softmax internally.

Try it yourself#

Run replace_head on models.swin_t() and models.regnet_y_400mf() first. You should get head with 768 features and fc with 440 features — two more names, two more widths.

Then run it on models.segmentation.fcn_resnet50(weights=None, weights_backbone=None):

IndexError: list index out of range

There is no nn.Linear anywhere in that model. Segmentation heads are convolutional, because the output is a map rather than a row of scores. Making replace_head handle a Conv2d head is the exercise.

What to learn next#

• Linear probing versus fine-tuning — training only the head, measured against training everything.
• Module registration and parameters — why assigning a new layer is enough for PyTorch to notice.
• Vision transformers — the family whose head lives at heads.head.