Repository: huggingface/transformers

Issue:
Missing `missing_keys` when loading from saved base model checkpoint
# 🐛 Bug

## Information

If a base model (e.g. `BertModel`, `DistilBertModel`, ...) is saved using `save_pretrained` and a model with an additional head (e.g. `BertForSequenceClassification`, `DistilBertForQuestionAnswering`, ...) is loaded from that checkpoint, it will not detect that it is missing layers.

## To reproduce

Steps to reproduce the behavior:

1. Instantiate base model from configuration or from `from_pretrained`
2. Save model using `save_pretrained`
3. Load checkpoint in model with head
4. No warning is output. Furthermore, if `output_loading_info=True` in step 3), will output `{'missing_keys': [], 'unexpected_keys': [], 'error_msgs': []}`

Here's a reproducible example:

```py
from transformers import BertForSequenceClassification, BertModel, BertConfig

config = BertConfig()
base_model = BertModel(config)
base_model.save_pretrained(directory)

model, loading_info = BertForSequenceClassification.from_pretrained(directory, output_loading_info=True)
print(loading_info)

# {'missing_keys': [], 'unexpected_keys': [], 'error_msgs': []}
# Should output {'missing_keys': ['classifier.weight', 'classifier.bias'], 'unexpected_keys': [], 'error_msgs': []}

```

## Expected behavior

Should detect the missing keys, as it does when loading from a full checkpoint:

```py
from transformers import BertForSequenceClassification

model, loading_info = BertForSequenceClassification.from_pretrained("bert-base-cased", output_loading_info=True)
print(loading_info)

# {'missing_keys': ['classifier.weight', 'classifier.bias'], 'unexpected_keys': ['cls.predictions.bias', 'cls.predictions.transform.dense.weight', 'cls.predictions.transform.dense.bias', 'cls.predictions.decoder.weight', 'cls.seq_relationship.weight', 'cls.seq_relationship.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.transform.LayerNorm.bias'], 'error_msgs': []}

```

## Environment info
<!-- You can run the command `python transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->
     
- `transformers` version: master branch
- Platform: Linux-5.5.7-arch1-1-x86_64-with-arch
- Python version: 3.6.10
- PyTorch version (GPU?): 1.4.0 (True)
- Tensorflow version (GPU?): 2.1.0 (True)
- Using GPU in script?: No
- Using distributed or parallel set-up in script?: No



Discussion:

Resolving patch:
diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
--- a/src/transformers/modeling_utils.py
+++ b/src/transformers/modeling_utils.py
@@ -540,6 +540,15 @@ def load(module: nn.Module, prefix=""):
                 model_to_load = getattr(model, cls.base_model_prefix)
 
             load(model_to_load, prefix=start_prefix)
+
+            if model.__class__.__name__ != model_to_load.__class__.__name__:
+                base_model_state_dict = model_to_load.state_dict().keys()
+                head_model_state_dict_without_base_prefix = [
+                    key.split(cls.base_model_prefix + ".")[-1] for key in model.state_dict().keys()
+                ]
+
+                missing_keys.extend(head_model_state_dict_without_base_prefix - base_model_state_dict)
+
             if len(missing_keys) > 0:
                 logger.info(
                     "Weights of {} not initialized from pretrained model: {}".format(


Test patch:
diff --git a/tests/test_modeling_common.py b/tests/test_modeling_common.py
--- a/tests/test_modeling_common.py
+++ b/tests/test_modeling_common.py
@@ -526,6 +526,21 @@ def test_model_common_attributes(self):
             x = model.get_output_embeddings()
             self.assertTrue(x is None or isinstance(x, torch.nn.Linear))
 
+    def test_correct_missing_keys(self):
+        config, _ = self.model_tester.prepare_config_and_inputs_for_common()
+
+        for model_class in self.all_model_classes:
+            model = model_class(config)
+            base_model_prefix = model.base_model_prefix
+
+            if hasattr(model, base_model_prefix):
+                with tempfile.TemporaryDirectory() as temp_dir_name:
+                    model.base_model.save_pretrained(temp_dir_name)
+                    model, loading_info = model_class.from_pretrained(temp_dir_name, output_loading_info=True)
+
+                    with self.subTest(msg="Missing keys for {}".format(model.__class__.__name__)):
+                        self.assertGreater(len(loading_info["missing_keys"]), 0)
+
     def test_tie_model_weights(self):
         if not self.test_torchscript:
             return