Saving/Loading models in AllenNLP package

user_1716733 0
Viewed 2318

I am trying to load an AllenNLP model weights. I could not find any documentation on how to save/load a whole model, so playing with weights only.

from allennlp.nn import util
model_state = torch.load(filename_model, map_location=util.device_mapping(-1))
model.load_state_dict(model_state)

I modified my input corpus a bit and I am guessing because of this I am getting corpus-size mismatch:

RuntimeError: Error(s) in loading state_dict for BasicTextFieldEmbedder:

    size mismatch for token_embedder_tokens.weight: 
    copying a param with shape torch.Size([2117, 16]) from checkpoint, 
    the shape in current model is torch.Size([2129, 16]).

Seemingly there is no official way to save model with corpus vocabulary. Any hacks around it?

allennlp pytorch
1 Answers

There's a functionality in AllenNLP allowing to load or save a model. Have you followed the steps outlined in AllenNLP's tutorial? Below I pasted a snippet from the tutorial that might be of interest to you:

# Here's how to save the model.
with open("/tmp/model.th", 'wb') as f:
    torch.save(model.state_dict(), f)

vocab.save_to_files("/tmp/vocabulary")

# And here's how to reload the model.
vocab2 = Vocabulary.from_files("/tmp/vocabulary")

model2 = LstmTagger(word_embeddings, lstm, vocab2)
with open("/tmp/model.th", 'rb') as f:
    model2.load_state_dict(torch.load(f))

If the above somehow doesn't work for you, you can check allennlp.models.archival.archive_model helper function. Using this function you should be able to archive your model's training configuration along with weights and vocabulary to model.tar.gz. Here you can find more information on the constraints of the two approaches that I discussed

user_122224690
Related