Can't Load ViT Model for Fine Tuning

nielsr
Wed, 10 Aug 2022 09:33:19 +0000

Hi,

The reason this works for ‘google/vit-base-patch16-224-in21k’ but not for checkpoints like ‘google/vit-base-patch16-224’ is because the latter include a fine-tuned head on top (namely, a head with 1000 output neurons, as these were fine-tuned on ImageNet-1k).

However, as you’d like to use this model but change the number of output neurons to 5, you need to add the additional ignore_mismatched_sizes=True argument. This ensures that the fine-tuned with 1000 output neurons is replaced by a randomly initialized head with 5 output neurons.

model = TFViTForImageClassification.from_pretrained(
    model_name_or_path,
    num_labels=len(labels),
    id2label={str(i): c for i, c in enumerate(labels)},
    label2id={c: str(i) for i, c in enumerate(labels)},
    ignore_mismatched_sizes=True, # add this to replace the head
)