[solved] KeyError: 'unexpected key "module.encoder.embedding.weight" in state_dict'

wasiahmad: <p>I am getting the following error while trying to load a saved model.</p>
<p><code>KeyError: 'unexpected key "module.encoder.embedding.weight" in state_dict'</code></p>
<p>This is the function I am using to load a saved model. (as suggested in <a href="http://pytorch.org/docs/notes/serialization.html#recommend-saving-models" rel="nofollow noopener">http://pytorch.org/docs/notes/serialization.html#recommend-saving-models</a>)</p>
<pre><code>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))
</code></pre>
<p>The model is a sequence-to-sequence network whose init function (constructor) is given below.</p>
<pre><code>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)
</code></pre>
<p>When I print the model (sequence-to-sequence network), I get the following.</p>
<pre><code>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 -&gt; 12)
    (attn_combine): Linear (600 -&gt; 300)
    (drop): Dropout (p = 0.25)
    (out): Linear (300 -&gt; 43723)
    (rnn): LSTM(300, 300, batch_first=True, dropout=0.25)
  )
  (criterion): NLLLoss (
  )
)
</code></pre>
<p>So, <code>module.encoder.embedding</code> is an embedding layer, and <code>module.encoder.embedding.weight</code> represents the associated weight matrix. So, why it says- <code>unexpected key "module.encoder.embedding.weight" in state_dict</code>?</p>

fmassa: <p>You probably saved the model using <code>nn.DataParallel</code>, which stores the model in <code>module</code>, and now you are trying to load it without <code>DataParallel</code>. You can either add a <code>nn.DataParallel</code> temporarily in your network for loading purposes, or you can  load the weights file, create a new ordered dict without the <code>module</code> prefix, and load it back.</p>

wasiahmad: <p>Yes, I used <code>nn.DataParallel</code>. I didn’t understand your second suggestion. Loading the weights file, create a new ordered dict without the module prefix and load it back. (Can you provide an example?)</p>
<p>Are you suggesting something like this? (example taken from here - <a href="https://github.com/OpenNMT/OpenNMT-py/blob/master/train.py" rel="nofollow noopener">https://github.com/OpenNMT/OpenNMT-py/blob/master/train.py</a>)</p>
<pre><code>model_state_dict = model.module.state_dict() if len(opt.gpus) &gt; 1 else model.state_dict()
model_state_dict = {k: v for k, v in model_state_dict.items() if 'generator' not in k}
generator_state_dict = model.generator.module.state_dict() if len(opt.gpus) &gt; 1 else model.generator.state_dict()
#  (4) drop a checkpoint
checkpoint = {
	'model': model_state_dict,
	'generator': generator_state_dict,
	'dicts': dataset['dicts'],
	'opt': opt,
	'epoch': epoch,
	'optim': optim
}
torch.save(checkpoint,
		   '%s_acc_%.2f_ppl_%.2f_e%d.pt' % (opt.save_model, 100*valid_acc, valid_ppl, epoch))  
</code></pre>
<p>May I ask you one question about the above code snippet, what is <code>generator</code> here?</p>

fmassa: <p>I was thinking about something like the following:</p>
<pre><code class="lang-python"># 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)
</code></pre>

wasiahmad: <p>First of all, thanks a lot, by adding a <code>nn.DataParallel</code> temporarily in my network for loading purposes worked. Even I tried your second suggested approach, it worked for me as well. Thanks a lottt <img src="//discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=5" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

achaiah: <p><a class="mention" href="/u/wasiahmad">@wasiahmad</a>   By adding <code>nn.DataParallel</code> temporarily into your network did you have to have the same number of GPUs available to load the model as when you saved the model?</p>

Tingfan_Wu: <p>A related question, given the fact we see that saving DataParallel wrapped model can cause problems when the model_state_dict is loaded into an unwrapped model. Would one recommend to save the “unwrapped” ‘module’ field inside a DataParallel instance instead ?<br>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/parallel/data_parallel.py#L49" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/parallel/data_parallel.py#L49" target="_blank" rel="nofollow noopener">pytorch/pytorch/blob/master/torch/nn/parallel/data_parallel.py#L49</a></h4>
<pre class="onebox"><code class="lang-py"><ol class="start lines" start="39" style="counter-reset: li-counter 38 ;">
<li>
</li>
<li># TODO: update notes/cuda.rst when this class handles 8+ GPUs well</li>
<li>
</li>
<li>def __init__(self, module, device_ids=None, output_device=None, dim=0):</li>
<li>    super(DataParallel, self).__init__()</li>
<li>
</li>
<li>    if not torch.cuda.is_available():</li>
<li>        self.module = module</li>
<li>        self.device_ids = []</li>
<li>        return</li>
<li class="selected">
</li>
<li>    if device_ids is None:</li>
<li>        device_ids = list(range(torch.cuda.device_count()))</li>
<li>    if output_device is None:</li>
<li>        output_device = device_ids[0]</li>
<li>    self.dim = dim</li>
<li>    self.module = module</li>
<li>    self.device_ids = device_ids</li>
<li>    self.output_device = output_device</li>
<li>    if len(self.device_ids) == 1:</li>
<li>        self.module.cuda(device_ids[0])</li>
</ol></code></pre>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>
</p>

Eugenio_Culurciello: <p>here is our way for alexnet trained with pytorch examples imagenet:</p><aside class="quote" data-post="15" data-topic="349">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/eugenio_culurciello/48/52_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/on-a-cpu-device-how-to-load-checkpoint-saved-on-gpu-device/349/15">On a cpu device, how to load checkpoint saved on gpu device</a> 
  </div>
  <blockquote>
    We have trained an Alexnet with pytorch examples imagenet (<a href="https://github.com/pytorch/examples/blob/master/imagenet/main.py" rel="noopener nofollow ugc">https://github.com/pytorch/examples/blob/master/imagenet/main.py</a>) and have been struggling to convert the model for use on CPU and for inference only. Here is a solution for AlexNet: 
<a href="https://github.com/e-lab/pytorch-toolbox/blob/master/convert-save-load.md" rel="noopener nofollow ugc">https://github.com/e-lab/pytorch-toolbox/blob/master/convert-save-load.md</a> 
It would be nice to have something more generic…
  </blockquote>
</aside>


Daniel_Du: <p>This works for me. Thanks a lot !</p>

datdddd: <p>I am having the same problem, and using the trick with OrderedDict does not work. I am using pytorch 0.3 in the case anything has changed.</p>
<p>I have an word embedding layer that was trained along with the classification task. Training was successful, but loading the model gave the error</p>
<p><code>Traceback (most recent call last):   File "source/test.py", line 72, in &lt;module&gt;     helper.load_model_states_from_checkpoint(model, args.save_path + 'model_best.pth.tar', 'state_dict', args.cuda)   File "/u/flashscratch/flashscratch1/d/datduong/universalSentenceEncoder/source/helper.py", line 55, in load_model_states_from_checkpoint     model.load_state_dict(checkpoint[tag])   File "/u/home/d/datduong/project/anaconda3/lib/python3.5/site-packages/torch/nn/modules/module.py", line 490, in load_state_dict     .format(name)) KeyError: 'unexpected key "embedding.embedding.embedding.weight" in state_dict'</code></p>
<p>The key embedding.embedding.embedding.weight exists (see image). Please let me know what to do.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/2X/6/6453bb757041e5dc6d09ba143a3822d90bbe3455.PNG" data-download-href="https://discuss.pytorch.org/uploads/default/6453bb757041e5dc6d09ba143a3822d90bbe3455" title="Capture2"><img src="https://discuss.pytorch.org/uploads/default/optimized/2X/6/6453bb757041e5dc6d09ba143a3822d90bbe3455_2_690x301.PNG" alt="Capture2" data-base62-sha1="ejx8Q6gC1eUdAICRpZaXVeodrfL" width="690" height="301" srcset="https://discuss.pytorch.org/uploads/default/optimized/2X/6/6453bb757041e5dc6d09ba143a3822d90bbe3455_2_690x301.PNG, https://discuss.pytorch.org/uploads/default/optimized/2X/6/6453bb757041e5dc6d09ba143a3822d90bbe3455_2_1035x451.PNG 1.5x, https://discuss.pytorch.org/uploads/default/optimized/2X/6/6453bb757041e5dc6d09ba143a3822d90bbe3455_2_1380x602.PNG 2x" data-dominant-color="193E45"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Capture2</span><span class="informations">1448×632 107 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

Oktai15: <aside class="quote no-group quote-modified" data-username="fmassa" data-post="4" data-topic="1686">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/fmassa/48/58_2.png" class="avatar"> fmassa:</div>
<blockquote>
<h1>original saved file with DataParallel</h1>
<p>state_dict = torch.load(‘myfile.pth.tar’)</p>
<h1>create new OrderedDict that does not contain <code>module.</code></h1>
<p>from collections import OrderedDict<br>
new_state_dict = OrderedDict()<br>
for k, v in state_dict.items():<br>
name = k[7:] # remove <code>module.</code><br>
new_state_dict[name] = v</p>
<h1>load params</h1>
<p>model.load_state_dict(new_state_dict)</p>
</blockquote>
</aside>
<p>In my opinion, this question-answer should be in something FAQ <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

Wyn_Mew: <p>Check out your saved model file:</p>
<pre><code class="lang-python">check_point = torch.load('myfile.pth.tar')
check_point.key()
</code></pre>
<p>You may find out your ‘check_point’ got several keys such as ‘state_dict’ etc.</p>
<pre><code class="lang-python">checkpoint = torch.load(resume)
state_dict =checkpoint['state_dict']

from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in state_dict.items():
    name = k[7:] # remove 'module.' of dataparallel
    new_state_dict[name]=v

model.load_state_dict(new_state_dict)
</code></pre>

dashesy: <p>What about <code>nn.DistributedDataParallel</code>, it seems <code>DistributedDataParallel</code> and <code>DataParallel</code> can load each other’s parameters.<br>
Is there an official way to save/load among <code>DDP/DP/None</code>?</p>

17764591637: <p>just do this:</p>
<p>model = torch.load(train_model)<br>
…<br>
net.load_state_dict(model[‘state_dict’])</p>
<p>it works for me!</p>

harryfyodor: <p>Thanks a lot, this worked for me.</p>

Alon: <p>Instead of deleting the “module.” string from all the state_dict keys, you can save your model with:<br>
<code>torch.save(model.module.state_dict(), path_to_file)</code><br>
instead of<br>
<code>torch.save(model.state_dict(), path_to_file)</code><br>
that way you don’t get the “module.” string to begin with…</p>

stoneyang: <p>Thanks for your hints! It saved <a href="https://discuss.pytorch.org/t/missing-keys-unexpected-keys-in-state-dict-when-loading-self-trained-model/22379/2">my time</a><img src="https://discuss.pytorch.org/images/emoji/apple/rose.png?v=5" title=":rose:" class="emoji" alt=":rose:"></p>

amy_hu: <p>that’s work simple and perfect for me! thanks</p>

KaiyangZhou: <p>In case someone needs, this <a href="https://github.com/KaiyangZhou/deep-person-reid/blob/master/torchreid/utils/torchtools.py#L104" rel="nofollow noopener">function</a> can handle loading weights w/ and w/o ‘module’.</p>
<p>To save model without ‘module’, you may try <a href="https://github.com/KaiyangZhou/deep-person-reid/blob/master/torchreid/utils/iotools.py#L41" rel="nofollow noopener">this</a>.</p>

weblucas: <p>After pytorch 1.xx<br>
this was fixed, now you only need to do this</p>
<pre><code class="lang-auto">            if isinstance(args.pretrained, torch.nn.DataParallel):
                args.pretrained = args.pretrained.module
</code></pre>