RuntimeError: Error(s) in loading state_dict for Wav2Vec2ForCTC
Author: YingLi001
Pull request: no
Labels: name: bug

### System Info

Transformers version: 4.4.0
Platform: Google Colab
Python version: 3.7

### Who can help?

@patrickvonplaten, @anton-l, @sanchit-gandhi

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

Pre-trained model: "facebook/wav2vec2-large-xlsr-53"

I have fine-tuned the above pre-tained model on Timit dataset.

When I loaded my own dataset(named: Torgo), and try to evaluate the fine-tuned model performance on Torgo dataset, I got the following error:

`RuntimeError: Error(s) in loading state_dict for Wav2Vec2ForCTC:
	size mismatch for lm_head.weight: copying a param with shape torch.Size([64, 1024]) from checkpoint, the shape in current model is torch.Size([51, 1024]).
	size mismatch for lm_head.bias: copying a param with shape torch.Size([64]) from checkpoint, the shape in current model is torch.Size([51]).`


### Expected behavior

My expected behavior is that I can directly evaluate the fine-tuned model on Torgo dataset. In other words, how do I further train a fine-tuned model on a different dataset? 

After reading some materials on HuggingFace and github, I know it is due to the `config.vocab_size` of Torgo dataset is not match that of Timit dataset. 

My questions are as follows:

1. When I align the vocab size of Torgo dataset to the one of the fine-tuned model, do I need to guarantee the vacab extracted from Torgo dataset is as same as that extracted from Timit dataset? Or I just need to care about the size?

2. If I cannot align the vocab size to the one of the fine-tuned model, is there any other method that I can achieve the expected behavior?

3. Could I just load the encoder or decoder components from the fine-tuned model?

Because I am a beginner of the wav2vec2 model, if something that I mentioned above is wrong, feel free to correct me. Thank you in advance and looking forward to hearing from you.


Comments:
Hey @YingLi001! Great question, and awesome that you found help using materials on the HuggingFace Hub and GitHub!

I'll first provide some context regarding Wav2Vec2 models, their tokenisers and how they affects the model weights. This information should help in answering your questions!

The pre-trained Wav2Vec2 model maps a sequence of audio inputs to a sequence of hidden-state representations. In order to decode text form this, we need to map the hidden-state representations to a vector over our vocabulary. To do this, we add a linear layer on top of the pre-trained Wav2Vec2 model. This linear layer performs a linear transformation of our hidden-states. It maps them from a dimensionality of 1024 to a dimensionality equal to our vocabulary size. In the case of TIMIT, where we have a vocabulary size of 51, we map the hidden-state representations from 1024-d down to 51-d. To decode a single character from this, we'd simply take the argmax of the 51-d vector, and look-up the corresponding token from our tokenizer! So if we had a model that predicted the following 51-d vector:
```
[ 0.01 ]
[ 0.02 ]
[ 0.90 ]
...
[ 0.01 ]
```
The argmax would be the third token. The character that we'd predict would be the token at position 3 in the tokenizer. If we want to decode a string of character, we have to do something a bit more fancy than just taking the argmax (i.e. connectionist temporal classification (CTC)), but the linear transformation remains the same!

The linear layer that I've eluded to is called the "language model head", or `lm_head` for short. What's special about the LM head is that it has a dimensionality specific to the vocabulary that we train the model on. If we have a vocabulary of 51 tokens, we'll have an LM-head weight matrix of size [51, 1024] (map the 1024-d hidden-states to the 51-d output vector). If we have 64 tokens, such as in the Torgo dataset, we'll have an LM-head weight matrix of size [64, 1024].

Whenever you give a model a different vocabulary size, the LM-head is going to have to be reset to a new size. Because of this, the LM-head weights are going to be randomly initialised, and thus require fine-tuning if we want our model to generate sensible predictions. Since each dataset typically has a different vocabulary, we usually build a new tokeniser for each dataset, and fine-tune the Wav2Vec2 model accordingly.

1. If you match the vocabularies sizes one-to-one, it is possible to load the LM-head weights. However, this does not guarantee that your model will predict characters accurately. Suppose in TIMIT you built the following tokeniser:

```
"a": 1
"b": 2
"c": 3
...
"z": 26
```

And for Torgo you built a tokenizer of the same dimensionality, but with a re-ordered vocabulary:

```
"z": 1
"y": 2
"x": 3
...
"a": 26
```

If we now use our LM-head weights to make predictions, we might get the following vector:
```
[ 0.01 ]
[ 0.02 ]
[ 0.90 ]
...
[ 0.01 ]
```
Taking the argmax, we get the third token in our vocabulary. So for TIMIT, we'd output a "c". For Torgo, we'd output a "x". Very different! Because the vocabulary is shuffled, we've effectively re-initialised our LM head. If we want to load the LM-head and evaluate the model **without any further fine-tuning**, we would need to match the tokenisers **exactly** in vocabulary size and positions. This means that for the Torgo datasets, you would load the tokeniser that you built when you fine-tuned on TIMIT. However, if we permit fine-tuning on the Torgo dataset, we can do something a bit different. See point 3.

2. As mentioned, you need to align both the vocabulary size and the tokeniser exactly.

3. If the Torgo dataset is similar to TIMIT, it's more than valid to load the encoder in isolation from the TIMIT checkpoint and then train the model on Torgo. You could then build a tokenizer specifically for the Torgo dataset, but leverage the majority of the weights from the TIMIT checkpoint. You can do this with the `from_pretrained()` method specifying the checkpoint location, setting the `config.vocab_size` to the correct value, and `ingore_mismatched_sizes` to `True` (ignores the `RuntimeError` you got previously):
```python
config.vocab_size = VOCAB_SIZE
model = Wav2Vec2ForCTC.from_pretrained(CKPT_LOCATION, config=config, ingore_mismatched_sizes=True)
```
This will randomly initialise the LM-head weights. We can then go ahead and train the model on the Torgo dataset with our new purpose-built tokeniser to learn suitable weights.

Note that in total, we have approximately $51*1024 + 51 \approx 53 \times 10^{3}$ weights for the linear layer. Overall, the model has nearly 400M params. That means that we're only randomly initialising the last 0.01% of the model weights! The remaining 99.99% are loaded from the pre-trained checkpoint. This means that we need relatively little data to fine-tune a model when we randomly initialised the LM-head but load the rest of the weights from pre-trained.

What checkpoint you use to load your model before training on Torgo is at your discretion. If the datasets are similar, you could use the TIMIT checkpoint. Otherwise, you can fine-tune from scratch using the official pre-trained `facebook/wav2vec2-large-xlsr-53` checkpoint. In both cases, you'll likely have to build a new tokeniser to match the vocabulary of Torgo and randomly initialise the linear LM-head accordingly.

Hope that helps and best of luck with your task!

Hi @sanchit-gandhi 

Thank you sooo much for your quickly reply and providing the detailed and constructive feedback!

I followed your suggestions(point 3), and successfully load the fine-tuned checkpoint. However, I still encountered some problems. Could you provide some help? Thank you in advance. The code of load the fine-tuned checkpoint shows as follows:

`    config = Wav2Vec2Config.from_json_file(
        "/content/drive/MyDrive/Thesispackage/wav2vec2-base-timit-demo-phones/checkpoint-11500/config.json")

    config.vocab_size = VOCAB_SIZE
    print("vocab_size", config.vocab_size)
    
    model = Wav2Vec2ForCTC.from_pretrained(
        "/content/drive/MyDrive/Thesispackage/wav2vec2-base-timit-demo-phones/checkpoint-11500", 
        config=config, 
        ignore_mismatched_sizes=True,
        # ctc_loss_reduction="mean",
        # pad_token_id=processor.tokenizer.pad_token_id,
        # vocab_size=len(processor.tokenizer),
    )
    model.gradient_checkpointing_enable()
    model.pad_token_id=processor.tokenizer.pad_token_id
    model.ctc_loss_reduction="mean"
    model.vocab_size=len(processor.tokenizer)`



Q1: When I load the fine-tuned checkpoint from local path using the `Wav2Vec2ForCTC.from_pretrained()`, I cannot add the `ctc_loss_reduction`, `pad_token_id`, `vocab_size` attributes inside this method. It will give me an error **TypeError: __init__() got an unexpected keyword argument 'vocab_size' site**. Therefore, I add them after loading the model. Am I right? 

Q2: Because the `vocab_size` between Torgo and Timit is still different, **[** Actually, **Torgo** dataset has **51** phonemes. **Timit** dataset has **64** phonemes. The 51 phonemes of Torgo dataset are **included** in the 64 phonemes of Timit dataset. **]** when I fine-tuned the fine-tuned checkpoint of Timit dataset on Torgo dataset, I got the following error. 

`[/usr/local/lib/python3.7/dist-packages/torch/nn/functional.py](https://localhost:8080/#) in ctc_loss(log_probs, targets, input_lengths, target_lengths, blank, reduction, zero_infinity)
   2615         )
   2616     return torch.ctc_loss(
-> 2617         log_probs, targets, input_lengths, target_lengths, blank, _Reduction.get_enum(reduction), zero_infinity
   2618     )
   2619 

RuntimeError: blank must be in label range`

After looking for some materials on HuggingFace, I found this [link](https://discuss.huggingface.co/t/runtimeerror-blank-must-be-in-label-range/4976/5). But I did not find a corresponding answer. 

Do you have any suggestions for solving this problem? 

Thanks again and looking forward to hearing from you.





Hey @YingLi001, sorry for the late reply!

A1: The way you've set the vocab size with the config is entirely correct!

A2: Interesting! I would then re-use the tokenizer for TIMIT dataset when fine-tuning on Torgo and load the Wav2Vec2 checkpoint in it's entirety. There seems to be no need in building a new tokenizer if the vocabularies overlap entirely. You'll then also retain the knowledge of the last linear layer (LM head) when you load the checkpoint.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

Hey @YingLi001! I hope the above comments answered your questions. Feel free to reopen this issue if you're still encountering problems loading the `state_dict`, or a new issue if there's something else you're having issues with! More than happy to help 🤗