Adding an argument to exclude some states (pretrained weights) from being loaded.
Author: qqaatw
Pull request: no

# 🚀 Feature request

Adding an argument in `from_pretrained` to exclude some states (pretrained weights) from being loaded.

## Motivation

In general, we usually use `from_pretrained` method to load pretrained states, from CDN or local files, into the model. However, In case when I need to adjust the shape of certain layers (submodules), the errors will be raised due to mismatched shapes.

For example, in the following snippets, I changed the embedding_size of Electra in order to tie the same embeddings as BERT in the subsequent code, but due to the mismatched shapes, many RuntimeErrors were raised in `module._load_from_state_dict`.

```
from transformers import BertModel, BertConfig, ElectraModel, ElectraConfig

bert_config = BertConfig.from_pretrained('bert-base-uncased')
bert_model = BertModel.from_pretrained('bert-base-uncased')

electra_config = ElectraConfig.from_pretrained(
    'google/electra-small-generator',
    embedding_size=bert_config.hidden_size
)
electra_model = ElectraModel.from_pretrained('google/electra-small-generator', config=electra_config)
```

```
Exception has occurred: RuntimeError
Error(s) in loading state_dict for ElectraModel:
    size mismatch for electra.embeddings.word_embeddings.weight: copying a param with shape torch.Size([30522, 128]) from checkpoint, the shape in current model is torch.Size([30522, 768]).
    size mismatch for electra.embeddings.position_embeddings.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 768]).
    size mismatch for electra.embeddings.token_type_embeddings.weight: copying a param with shape torch.Size([2, 128]) from checkpoint, the shape in current model is torch.Size([2, 768]).
    size mismatch for electra.embeddings.LayerNorm.weight: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([768]).
    size mismatch for electra.embeddings.LayerNorm.bias: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([768]).
    size mismatch for electra.embeddings_project.weight: copying a param with shape torch.Size([256, 128]) from checkpoint, the shape in current model is torch.Size([256, 768]).
```

Therefore, I think it would be better to add an argument like `excluded_keys` (as the following example) in `from_pretrained` to explicitly prevent certain states from being loaded or add an argument to automatically have the states with mismatched shapes not loaded. I know there are some workarounds such as loading all states first then tying each weight respectively, but that will result in a long and not concise code segment.

Example:
```
electra_model = ElectraModel.from_pretrained(
    'google/electra-small-generator',
    config=electra_config,
    excluded_keys = [
        "electra.embeddings.word_embeddings.weight",
        "electra.embeddings.position_embeddings.weight",
        "electra.embeddings.token_type_embeddings.weight",
        "electra.embeddings.LayerNorm.weight",
        "electra.embeddings.LayerNorm.bias",
        "electra.embeddings_project.weight",
        "generator_predictions.LayerNorm.weight",
        "generator_predictions.LayerNorm.bias",
        "generator_predictions.dense.weight",
        "generator_predictions.dense.bias",
        "generator_lm_head.weight"
    ]
)
```

## Your contribution

If there is no other concern, and no one is implementing similar features, I would be happy to submit a PR for this.

Any thoughts are welcomed :)

Comments:
I'm in favor of this. For example, when I wanted to fine-tune DETR, I had to use the following hack to make it work:

```
model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
state_dict = model.state_dict()
# Remove class weights
del state_dict["class_labels_classifier.weight"]
del state_dict["class_labels_classifier.bias"]
# define new model with custom class classifier
config = DetrConfig.from_pretrained("facebook/detr-resnet-50", num_labels=10)
model = DetrForObjectDetection(config)
model.load_state_dict(state_dict, strict=False)
```

This is because DETR has a head that has been fine-tuned on COCO, and it has 91 classes. However, when fine-tuning on my custom dataset, let's say it has 10 labels, then the classification head needs to be replaced, which is what I do above.

It would be easier if I could just do (with an API similar to what you propose above):

```
model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50", num_labels=10, excluded_keys = [
        "class_labels_classifier.weight",
        "class_labels_classifier.bias",
        ]
)
```

This way, you can easily replace the head of an already fine-tuned model with your custom classification head.

cc @patil-suraj @LysandreJik @patrickvonplaten @sgugger 

Interesting proposal! I would also be in favor of this to enable @qqaatw's use-case.

For your use-case @NielsRogge, while the proposal would also work, I'd favor something much simpler as it's more common to want to drop a head to load in the same architecture but with a randomly initialized layer. With the proposal here, it implies knowing the weight names and manually specifying them.

It can also be achieved with

```
model = DetrModel.from_pretrained("facebook/detr-resnet-50")
model.save_pretrained("directory")

model = DetrForObjectDetection.from_pretrained("directory")
```
which will randomly initialize all layers that are new in `DetrForObjectDetection`.

For your use-case in particular Niels, I would be in favor of having an API like the following:
```
model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50", load_head=False)
```
It would imply being aware of the head layers, which would probably be achieved by using a `model.get_classification_head` similar to the existing `model.get_output_embeddings` method.

We have already discussed the second use case internally, and I think we came to the conclusion @NielsRogge code should work with a warning (like the ones we get when the head is different because we load a checkpoint for a task on another task).

The other use case presented in this issue is also interesting. Do we really need to add a new argument for it? We could treat it the same way: when trying to load the weights and there is a size mismatch, just ignore the weights and put them in the warning raised by the `from_pretrained` method.

@LysandreJik I tried your first code block, however, `DetrForObjectDetection` has 2 heads (one for class labels, one for bounding boxes), and one typically only wants to randomly initialize the class labels classifier (and further train the bounding box regressor head). However, your code only works if you want to randomly initialize both heads (it prints the following warning):

```
Some weights of the model checkpoint at facebook/detr-resnet-50 were not used when initializing DetrModel: ['bbox_predictor.layers.1.bias', 'bbox_predictor.layers.0.weight', 'bbox_predictor.layers.0.bias', 'class_labels_classifier.bias', 'bbox_predictor.layers.1.weight', 'class_labels_classifier.weight', 'bbox_predictor.layers.2.weight', 'bbox_predictor.layers.2.bias']
- This IS expected if you are initializing DetrModel 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 DetrModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).
Some weights of DetrForObjectDetection were not initialized from the model checkpoint at directory and are newly initialized: ['bbox_predictor.layers.1.bias', 'bbox_predictor.layers.0.weight', 'bbox_predictor.layers.0.bias', 'class_labels_classifier.bias', 'bbox_predictor.layers.1.weight', 'class_labels_classifier.weight', 'bbox_predictor.layers.2.weight', 'bbox_predictor.layers.2.bias']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
```

Of course, I don't think there are many models in the library right now that have multiple heads, but DETR is one such model.




After discussing offline with @LysandreJik we will add a `ignore_mismatched_size` flag to `from_pretrained`. When activated, weights that don't have the right size will be ignored, which should cover both the use cases in this issue.

I will work on this today.

The feature has been implemented in #12664. Thanks @sgugger