How do I change the classification head of a model?

This is now possible (thanks to @sgugger) by passing in an additional argument called ignore_mismatched_sizes, which you can set to True.

If you have an already fine-tuned model with, let’s say 17 labels, and you want to replace the head with one that has 10 outputs, you can do it as follows:

from transformers import BertForTokenClassification

model_name = "vblagoje/bert-english-uncased-finetuned-pos"

model = BertForTokenClassification.from_pretrained(model_name, num_labels=10, ignore_mismatched_sizes=True)

This will print the following warning:

Some weights of BertForTokenClassification were not initialized from the model checkpoint at vblagoje/bert-english-uncased-finetuned-pos and are newly initialized because the shapes did not match:
- classifier.weight: found shape torch.Size([17, 768]) in the checkpoint and torch.Size([10, 768]) in the model instantiated
- classifier.bias: found shape torch.Size([17]) in the checkpoint and torch.Size([10]) in the model instantiated
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
23 Likes