KeyError: "unexpected key" module.encoder.embedding.weight "in state_dict"
When trying to load a saved model, the following error appears.
KeyError: 'unexpected key "module.encoder.embedding.weight" in state_dict'
This is the function I am using to load the saved model.
def load_model_states(model, tag):
"""Load a previously saved model states."""
filename = os.path.join(args.save_path, tag)
with open(filename, 'rb') as f:
model.load_state_dict(torch.load(f))
The model is a sequence-to-sequence network, whose init (constructor) function is shown below.
def __init__(self, dictionary, embedding_index, max_sent_length, args):
""""Constructor of the class."""
super(Sequence2Sequence, self).__init__()
self.dictionary = dictionary
self.embedding_index = embedding_index
self.config = args
self.encoder = Encoder(len(self.dictionary), self.config)
self.decoder = AttentionDecoder(len(self.dictionary), max_sent_length, self.config)
self.criterion = nn.NLLLoss() # Negative log-likelihood loss
# Initializing the weight parameters for the embedding layer in the encoder.
self.encoder.init_embedding_weights(self.dictionary, self.embedding_index, self.config.emsize)
When I print the model (sequence to sequence) I get the following.
Sequence2Sequence (
(encoder): Encoder (
(drop): Dropout (p = 0.25)
(embedding): Embedding(43723, 300)
(rnn): LSTM(300, 300, batch_first=True, dropout=0.25)
)
(decoder): AttentionDecoder (
(embedding): Embedding(43723, 300)
(attn): Linear (600 -> 12)
(attn_combine): Linear (600 -> 300)
(drop): Dropout (p = 0.25)
(out): Linear (300 -> 43723)
(rnn): LSTM(300, 300, batch_first=True, dropout=0.25)
)
(criterion): NLLLoss (
)
)
So, it module.encoder.embedding
is the embedding layer, and it module.encoder.embedding.weight
represents the associated weight matrix. So why does he say : unexpected key "module.encoder.embedding.weight" in state_dict
?
source to share
I solved the problem. In fact, I was saving the model with a nn.DataParallel
which stores the model in a module and then I tried to load it without DataParallel
. So I need to temporarily add nn.DataParallel
to my network to load, or I can load the weights file, create a new ordered dict without the module prefix and load it back.
The second workaround looks like this.
# original saved file with DataParallel
state_dict = torch.load('myfile.pth.tar')
# create new OrderedDict that does not contain `module.`
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k[7:] # remove `module.`
new_state_dict[name] = v
# load params
model.load_state_dict(new_state_dict)
source to share