How to load part of pre trained model?

albertxavier001: <p>How to load part of pre trained model?<br>
Will the unused parameters be auto deleted if I load the whole model but only use part of it?</p>

chenyuntc: <p>I’m afraid not</p>
<blockquote>
<p>The keys of <code>state_dict</code> must exactly match the keys returned by this module’s  <code>state_dict()</code> function.</p>
</blockquote>
<p>if name not in new model , it will raise  KeyError<br>
but I gusee this may work for you</p>
<pre data-code-wrap="python"><code class="lang-python">    def load_my_state_dict(self, state_dict):
 
        own_state = self.state_dict()
        for name, param in state_dict.items():
            if name not in own_state:
                 continue
            if isinstance(param, Parameter):
                # backwards compatibility for serialized parameters
                param = param.data
            own_state[name].copy_(param)
</code></pre>

apaszke: <p>You can remove all keys that don’t match your model from the state dict and use it to load the weights afterwards:</p>
<pre><code class="lang-auto">pretrained_dict = ...
model_dict = model.state_dict()

# 1. filter out unnecessary keys
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
# 2. overwrite entries in the existing state dict
model_dict.update(pretrained_dict) 
# 3. load the new state dict
model.load_state_dict(pretrained_dict)
</code></pre>

albertxavier001: <p>Wow, thanks a lot!<br>
But “model.load(pretrained_dict)”  gives me an error “object has no attribute ‘load’”</p>

albertxavier001: <p>I think  <span class="hashtag">#3</span> should be:</p>
<pre><code>model.load_state_dict(model_dict)</code></pre>

apaszke: <p>Yes it should. I edited the snippet.</p>

cham-3: <p>Thanks. I am looking for it. It seems awesome!!!</p>

ImgPrcSng: <p>Consider the situation where I would like to restore all weights till the last layer.</p>
<pre><code class="lang-auto"># args has the model name, num classes and other irrelevant stuff
self._model = models.__dict__[args.arch](pretrained = False, 
                                         num_classes = args.classes, 
                                         aux_logits = False)


if self.args.pretrained:
              
    print("=&gt; using pre-trained model '{}'".format(args.arch))
    pretrained_state = model_zoo.load_url(model_names[args.arch])
    model_state = self._model.state_dict()

    pretrained_state = { k:v for k,v in pretrained_state.iteritems() if k in model_state and v.size() == model_state[k].size() }
    model_state.update(pretrained_state)
    self._model.load_state_dict(model_state)
</code></pre>
<p>Shouldn’t we also be checking if the sizes match before restoring? It looks like we are comparing only the names.</p>

smth: <p>if the sizes are wrong, i believe the <code>copy_</code> invoked in <code>load_state_dict</code> will complain.</p>

BrianHuang: <h1>3. load the new state dict</h1>
<p>model.load_state_dict(model_dict)</p>
<p>step 3 should look like this</p>

chenchr: <p><a class="mention" href="/u/chenyuntc">@chenyuntc</a> Hello, what if the net2 is a subset of net1, and I want to load weight from net2 to net1? can directly using load_static_dict works? thanks!</p>

karttikeya_mangalam: <p>As of 21st December’17,  load_state_dict() takes the boolean argument ‘strict’ which when set to False allows to load only the variables that are identical between the two models irrespective of whether one is subset/superset of the other.<br>
<a href="http://pytorch.org/docs/master/_modules/torch/nn/modules/module.html#Module.load_state_dict" class="onebox" target="_blank" rel="nofollow noopener">http://pytorch.org/docs/master/_modules/torch/nn/modules/module.html#Module.load_state_dict</a></p>

zeakey: <p>After <code>model_dict.update(pretrained_dict)</code>, the <code>model_dict</code> may still have keys that <code>pretrained_model</code> doesn’t have, which will cause a error.</p>
<p>Assum following situation:</p>
<pre><code class="lang-auto">pretrained_dict: ['A', 'B', 'C', 'D']
model_dict: ['A', 'B', 'C', 'E']
</code></pre>
<p>After <code>pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}</code> and <code>model_dict.update(pretrained_dict)</code>, they are:</p>
<pre><code class="lang-auto">pretrained_dict: ['A', 'B', 'C']
model_dict: ['A', 'B', 'C', 'E']
</code></pre>
<p>So when performing <code>model.load_state_dict(pretrained_dict)</code>, <code>model_dict</code> still has key <code>E</code> that <code>pretrained_dict</code> doen’t have.</p>
<p>So how about using <code>model.load_state_dict(model_dict)</code> instead of <code>model.load_state_dict(pretrained_dict)</code>?</p>
<p>The complete snippet is therefore as follow:</p>
<pre><code class="lang-auto">pretrained_dict = ...
model_dict = model.state_dict()

# 1. filter out unnecessary keys
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
# 2. overwrite entries in the existing state dict
model_dict.update(pretrained_dict) 
# 3. load the new state dict
model.load_state_dict(model_dict)
</code></pre>

romina.ui: <p>hi<br>
how I can load trained model with its weights ?</p>

ptrblck: <p>Have a look at the <a href="http://pytorch.org/docs/stable/notes/serialization.html">Serialization tutorial</a> or do you have a specific issue?</p>

sriharsha0806: <p>Is it param = Parameter.data ?</p>

deJQK: <p>It might be a little late to ask. If I want to skip some layers, like if I train the model with batch normalization, but want to use the trained bn version for that without batch normalization, how can I change the layers’ names? Because otherwise, the name might be different, and it will complain about size mismatching.</p>

kingxueyuf: <p>This is very useful, thanks</p>

nurlano: <p>Shouldn’t the # 3  be</p>
<pre><code class="lang-auto">model.load_state_dict(model_dict)
</code></pre>
<p>instead of</p>
<pre><code class="lang-auto">model.load_state_dict(pretrained_dict)
</code></pre>
<p>?</p>

Zichun_Zhang: <p>if those params are stored in dictionary,<br>
there is a great a lot parameters that have same names, eg. nn.Conv2d, really a lot.<br>
So how could it know which nn.Conv2d is the right one to choose, when use this dictionary to mapping.<br>
Thx~</p>