RuntimeError: Error(s) in loading state_dict for BertModel

Question

I finetune a BERT model using hugging face transformer library and train it in GPU in the cloud. Then I save the model and tokenizer like below:

model.save_pretrained('/saved_model/')
torch.save(best_model.state_dict(), '/saved_model/model')
tokenizer.save_pretrained('/saved_model/')

I download the saved_model directory in my computer. Then I load the model/tokenizer like below in my computer

import torch
from transformers import *
tokenizer = BertTokenizer.from_pretrained('./saved_model/')
config = BertConfig('./saved_model/config.json')
model = BertModel(config)
model.load_state_dict(torch.load('./saved_model/pytorch_model.bin', map_location=torch.device('cpu')))
model.eval()

But it throws below error for the model.load_state_dict line

RuntimeError: Error(s) in loading state_dict for BertModel:
    Missing key(s) in state_dict:

It lists a bunch of keys that are apparently missing from the state_dict.

I am new to pytorch and not sure what is going on. Most likely I am not saving the model the right way.

Please suggest.

Wasi Ahmad Wasi Ahmad · Accepted Answer · 2019-10-18T08:06:48

As you may know, the state_dict of a PyTorch module is an OrderedDict. When you tried to load the weights of a module from a state_dict, it complains about missing keys which means the state_dict does not contain those keys. In this situation, I would suggest taking the following actions.

Check which keys are present in the state_dict. It sounds impossible that you save a subset of the keys only.
Also, make sure you have the correct configuration loaded. Otherwise, if your trained BertModel and the new BertModel for which you want to load the weights are different, then you will receive this error.
Finally, if your code passes both the above cases, then saving the model, make sure you save all the layers' parameters in the file. The statement, torch.save(best_model.state_dict(), '/saved_model/model') looks okay to me but make sure the best_model.state_dict() contains all the expected keys.

RuntimeError: Error(s) in loading state_dict for BertModel

1 Answers