# Loading Vision Transformer Model After Changing Its Classifier Head

**URL:** <https://discuss.huggingface.co/t/loading-vision-transformer-model-after-changing-its-classifier-head/55492>\
**Category:** 🤗Transformers\
**Created:** [September 18, 2023, 4:34pm UTC](https://discuss.huggingface.co/t/loading-vision-transformer-model-after-changing-its-classifier-head/55492 "2023-09-18T16:34:02Z")\
**Posts on this page:** 3\
**Page:** 1

<div class="post-metadata">

**Author:** ![mohie](https://avatars.discourse-cdn.com/v4/letter/m/d2c977/32.png) [@mohie](https://discuss.huggingface.co/u/mohie)\
**Post date:** [September 18, 2023, 4:34pm UTC](https://discuss.huggingface.co/t/loading-vision-transformer-model-after-changing-its-classifier-head/55492/1 "2023-09-18T16:34:02Z")

</div>

I added more layers in the classifier head of the VIT model for image classification task like this:

```auto
model = ViTForImageClassification.from_pretrained("google/vit-base-patch16-224-in21k")

# Define a new sequential module for the classifier
new_classifier = nn.Sequential(
    nn.Linear(768, 256), # First layer with 256 neurons
    nn.ReLU(), # Activation function
    nn.BatchNorm1d(256),# Batch normalization for the first layer
    nn.Dropout(0.1), # dropout
    nn.Linear(256, 64), # Second layer with 64 neurons
    nn.ReLU(), # Activation function
    nn.BatchNorm1d(64), # Batch normalization for the second layer
    nn.Dropout(0.1), # dropout
    nn.Linear(64, 2) # Output layer with 2 neurons
)
model.classifier = new_classifier

```

Then I trained the model using Trainer() class from hugging face and saved the best model in a checkpoint.  
when loading the trained model from this checkpoint I get this warning:

```auto
Some weights of the model checkpoint at /content/best_models_complex/checkpoint-564 were not used when initializing ViTForImageClassification: ['classifier.8.weight', 'classifier.6.running_var', 'classifier.2.running_mean', 'classifier.6.num_batches_tracked', 'classifier.6.weight', 'classifier.4.bias', 'classifier.2.running_var', 'classifier.2.num_batches_tracked', 'classifier.6.bias', 'classifier.8.bias', 'classifier.2.bias', 'classifier.2.weight', 'classifier.0.weight', 'classifier.0.bias', 'classifier.6.running_mean', 'classifier.4.weight']
- This IS expected if you are initializing ViTForImageClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).
- This IS NOT expected if you are initializing ViTForImageClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).
Some weights of ViTForImageClassification were not initialized from the model checkpoint at /content/best_models_complex/checkpoint-564 and are newly initialized: ['classifier.weight', 'classifier.bias']
You should probably TRAIN this model on a downstream task to be able to use it for predictions and inference.

```

And the predictions of the model are very bad (probably because of the random weights of the added layers)  
So my question is **how to load the model from the checkpoint with the weights of the new layers not being ranodm?**

---

<div class="post-metadata">

**Author:** ![mohie](https://avatars.discourse-cdn.com/v4/letter/m/d2c977/32.png) [@mohie](https://discuss.huggingface.co/u/mohie)\
**Post date:** [September 18, 2023, 4:36pm UTC](https://discuss.huggingface.co/t/loading-vision-transformer-model-after-changing-its-classifier-head/55492/2 "2023-09-18T16:36:20Z")

</div>

I also wonder how Trainer() can load the best model at the end of the training correctly if loading the model from a checkpoint is not working properly.

---

<div class="post-metadata">

**Author:** ![nielsr](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/nielsr/32/39617_2.png) [@nielsr](https://discuss.huggingface.co/u/nielsr)\
**Post date:** [December 21, 2023, 1:50pm UTC](https://discuss.huggingface.co/t/loading-vision-transformer-model-after-changing-its-classifier-head/55492/3 "2023-12-21T13:50:08Z")

</div>

Hi,

To define a custom ViT for image classification model, it should be defined like [here](https://github.com/huggingface/transformers/blob/f9a98c476c3a92beaba3b20e51f1ff49417231a6/src/transformers/models/vit/modeling_vit.py#L756), with self.classifier being replaced by your custom classification head. Important is that your class should inherit from ViTPreTrainedModel in order to make it work with `from_pretrained`. This class can be imported from transformers.models.vit.modeling\_vit.