How do you replace a pretrained model's 1000-class head for a 7-class task?

## Anatomy A classification network is a **backbone** that turns an input into a pooled feature vector — a few hundred to a few thousand numbers — followed by a **head**: a linear layer with a weight matrix of shape (number of classes) x (feature dimension), plus a bias per class. The head produces logits, which a softmax and a cross-entropy loss turn into a training signal. Each row of that weight matrix is, in effect, a prototype for one source class. Row 386 means whatever class 386 meant in the pretraining label set. Your seven classes have no correspondence to those rows, and even where a name coincidentally overlaps, the ordering and the calibration are different. There is nothing to salvage, so the head is replaced entirely rather than adapted. ## Building the replacement Construct a new linear layer with output width 7 and the same input width as the pooled representation. Initialise the weights small and random. Biases at zero is the standard default. When the classes are heavily imbalanced, initialising each bias to the log of that class's prior probability is a well-known improvement: with weights near zero, the model's first predictions are the base rates rather than a uniform guess, which avoids spending the first epochs merely learning the intercept and removes a large, uninformative initial loss spike. If the task is not single-label classification, the head changes in more than width. Multi-label means one independent output per label with a per-label binary cross-entropy rather than a softmax over classes; regression means a single linear output with a squared-error or absolute-error loss. Occasionally a small two-layer head helps, especially with a fully frozen backbone on a distant domain, at the cost of more parameters to overfit with a small dataset. ## Loading, shapes and mappings Loading a pretrained checkpoint strictly — every stored tensor must find a matching one — fails at the head, because the stored matrix has 1000 rows and yours has 7. That failure is expected: load the backbone's tensors and construct the head separately. The label mapping deserves more care than it usually gets. Your seven class names must map to indices 0-6 by a mapping that is written down, versioned, and used identically when training and when serving. A mapping that is derived from, say, directory listing order in one place and a hand-written list in another is a silent catastrophe: the model trains perfectly, the metrics look fine, and production predictions are systematically permuted. One more input-side check: the preprocessing must match what the backbone was pretrained with — the same resize convention and the same input scaling and channel normalization. Feed it differently scaled inputs and the pretrained features are being computed on out-of-distribution data, which quietly costs accuracy in a way no error message reveals. ## Prove the wiring on 32 examples Before any real run, take a fixed subset of 32 training examples, turn off augmentation and any regularisation, and train on that subset repeatedly. The training loss should fall to near zero and accuracy to 100% within a short run. This is not a model-quality result — it is a wiring check on a freshly rebuilt pipeline, and it is the cheapest one available. If it plateaus instead, look at the plumbing, not the data: - The new head's parameters are not in the set being optimised — an easy mistake when you froze the backbone by excluding parameters and swept up the head too. - The label mapping used in the loader disagrees with the label file. - The loss is receiving values in a form it does not expect — for example already-normalised probabilities where raw logits are required, or the wrong argument order. - Augmentation or dropout is still on, so the 32 examples are never actually seen twice. - Everything, head included, is frozen, so nothing is being learned at all. ## What comes next A randomly initialised head produces large, essentially arbitrary gradients in its first steps, and those gradients flow back into whatever part of the backbone you unfroze. How to keep that from disturbing good pretrained features is a step-size question with its own answer, but it is the reason the head swap and the training schedule are usually thought about together. The main thing to carry away: the backbone is what transfers, the head is a per-task attachment, and the boring parts around it — the label mapping, the preprocessing convention, and a 32-example proof that the pipeline learns at all — are where the real failures live.