How to use the helper function: consume prefix in state dict if_present

John_J_Watson: <p>he docs suggest to use:</p>
<pre><code class="lang-auto">To let a non-DDP model load a state dict from a DDP model, consume_prefix_in_state_dict_if_present() needs to be applied to strip the prefix “module.” in the DDP state dict before loading.
</code></pre>
<p>Source: <a href="https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html" class="inline-onebox" rel="noopener nofollow ugc">DistributedDataParallel — PyTorch 1.9.0 documentation</a></p>
<p>however., this line:</p>
<pre><code class="lang-auto">model_state_dict=torch.nn.modules.utils.consume_prefix_in_state_dict_if_present(checkpoint['model_dict'],prefix='module.')
</code></pre>
<p>gives back <code>None</code> <img src="https://discuss.pytorch.org/images/emoji/apple/frowning.png?v=10" title=":frowning:" class="emoji" alt=":frowning:"></p>
<p>I do something ugly like this at the moment:</p>
<pre><code class="lang-auto">        new_state_dict = collections.OrderedDict()
        for k, v in checkpoint['model_dict'].items():
            name = k.replace("module.", '') # remove `module.`
            new_state_dict[name] = v
</code></pre>
<p>but I would prefer to use the <code>consume_prefix_in_state_dict_if_present</code>.</p>
<p>Can someone elucidate the correct usage of this please? Obviously, I am not getting it!</p>

mrshenli: <p>calling for expert <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=10" title=":slight_smile:" class="emoji" alt=":slight_smile:">  cc <a class="mention" href="/u/wayi">@wayi</a></p>

wayi: <p>Hi <a class="mention" href="/u/john_j_watson">@John_J_Watson</a>, sorry for the confusion. <code>consume_prefix_in_state_dict_if_present</code> removes the prefix in place rather than returns any value. You just use <code>checkpoint['model_dict']</code> instead of creating a temporary variable here.</p>
<p>Example: <a href="https://github.com/pytorch/pytorch/blob/34c9f5a8dad74ba23de5c2fba9d071a6c2dd1fa4/test/distributed/test_c10d_gloo.py#L1818" class="inline-onebox" rel="noopener nofollow ugc">pytorch/test_c10d_gloo.py at 34c9f5a8dad74ba23de5c2fba9d071a6c2dd1fa4 · pytorch/pytorch · GitHub</a></p>
<p>I created a PR to improve the documentation: <a href="https://github.com/pytorch/pytorch/pull/63388" class="inline-onebox" rel="noopener nofollow ugc">Add return type hint and improve the docstring of consume_prefix_in_state_dict_if_present method by SciPioneer · Pull Req</a></p>

John_J_Watson: <p>thank you <a class="mention" href="/u/wayi">@wayi</a> for this. I dont know why I didnt think to try thinking it could be inplace!<br>
Just as a follow up quetsion, this wrapper basically does the same as follows?</p>
<pre><code class="lang-auto">        new_state_dict = collections.OrderedDict()
        for k, v in checkpoint['model_dict'].items():
            name = k.replace("module.", '') # remove `module.`
            new_state_dict[name] = v
</code></pre>
<p>Is that right? Are there any advantages/disadvantages of using either approaches?</p>
<p>Thank you again!</p>

wayi: <p>One subtle difference is that “_metadata” field (if any) is handled separately. See:</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/e000dfcf976454fdadfdc556248976e6e560d155/torch/nn/modules/utils.py#L63">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/e000dfcf976454fdadfdc556248976e6e560d155/torch/nn/modules/utils.py#L63" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/e000dfcf976454fdadfdc556248976e6e560d155/torch/nn/modules/utils.py#L63" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/e000dfcf976454fdadfdc556248976e6e560d155/torch/nn/modules/utils.py#L63</a></h4>



  <pre class="onebox">    <code class="lang-py">
      <ol class="start lines" start="53" style="counter-reset: li-counter 52 ;">
          <li>    state_dict (OrderedDict): a state-dict to be loaded to the model.</li>
          <li>    prefix (str): prefix.</li>
          <li>"""</li>
          <li>keys = sorted(state_dict.keys())</li>
          <li>for key in keys:</li>
          <li>    if key.startswith(prefix):</li>
          <li>        newkey = key[len(prefix) :]</li>
          <li>        state_dict[newkey] = state_dict.pop(key)</li>
          <li>
          </li>
<li># also strip the prefix in metadata if any.</li>
          <li class="selected">if "_metadata" in state_dict:</li>
          <li>    metadata = state_dict["_metadata"]</li>
          <li>    for key in list(metadata.keys()):</li>
          <li>        # for the metadata dict, the key can be:</li>
          <li>        # '': for the DDP module, which we want to remove.</li>
          <li>        # 'module': for the actual model.</li>
          <li>        # 'module.xx.xx': for the rest.</li>
          <li>
          </li>
<li>        if len(key) == 0:</li>
          <li>            continue</li>
          <li>        newkey = key[len(prefix) :]</li>
      </ol>
    </code>
  </pre>


  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>Other than that, I don’t think there is a big difference. Your own implementation is a little less memory efficient I will say, as you don’t do it in place.</p>