How to copy a modified state_dict into a models state_dict

username: Cedric_Gullentops
cooked: <p>I am trying to copy a modified state_dict from a model that was pruned (e.g. reducing the 0 dimension of one of the tensors by 1). The models have the same keys, the only difference is the dimensions of the tensors. I made a wrapper class to handle pruned networks as seen below:</p>
<pre><code class="lang-auto">class PrunedNetwork():
    def __init__(self, network):
        self.network = network

    def __getattr__(self, name):
        return getattr(self.network, name)

    def forward(self, *args, **kwargs):
        self.network.forward(*args, **kwargs)

    def save(self, filename):
        self.network.save(filename)

    def load(self, filename):
        new_state_dict = torch.load(filename, 'cpu')

        for param_tensor in self.network.state_dict():
            print(param_tensor, "\t", self.network.state_dict()[param_tensor].size())

        # TODO: copy new_state_dict into state_dict
        # The following attempt doesn't work
        keys = self.network.state_dict().keys()
        for key in keys:
            self.network.state_dict()[key] = new_state_dict[key]
        
        self.network.load(filename)
</code></pre>
<p>The tensors aren’t being copied over to the model.</p>

username: LeviViana
cooked: <p>I think you are only accessing a copy of the parameters data here:</p>
<aside class="quote no-group" data-username="Cedric_Gullentops" data-post="1" data-topic="64828">
<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/cedric_gullentops/48/18532_2.png" class="avatar"> Cedric_Gullentops:</div>
<blockquote>
<p>for key in keys: self.network.state_dict()[key] = new_state_dict[key]</p>
</blockquote>
</aside>
<p>Why not just this ?</p>
<pre data-code-wrap="python"><code class="lang-python">    def load(self, filename):
        new_state_dict = torch.load(filename, 'cpu')
        self.network.load_state_dict(new_state_dict)
</code></pre>
<p>EDIT:</p>
<p>The above won’t work if the dimensions don’t match. So, in your case, I would recreate all network from scratch but having the right dimensions, and simply do a <code>self.network.load_state_dict(...)</code>, like below:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch

net = torch.nn.Linear(2, 2)
new_state_dict = {'weight': torch.rand(1, 2), 'bias': torch.rand(1)}

net.load_state_dict(new_state_dict) # &lt;- this won't work

new_net = torch.nn.Linear(2, 1)
new_net.load_state_dict(new_state_dict) # &lt;- this works
</code></pre>

username: Cedric_Gullentops
cooked: <p><a class="mention" href="/u/leviviana">@LeviViana</a>  The problem is I currently have about 20 weight files of pruned networks (more on the way) so recreating every network by hand isn’t really doable. The goal is to read a network from a config file, edit the tensor dimensions to be the same as the pruned weight file and then load in the weights.</p>

username: Cedric_Gullentops
cooked: <p>Another thing I don’t understand is that</p>
<pre><code class="lang-auto">self.network.state_dict()[param_tensor][0][0] = 0
</code></pre>
<p>sets the values of this tensors dimension to 0 but assigning the full tensor doesn’t retain the changes.</p>

username: Cedric_Gullentops
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> <a class="mention" href="/u/alband">@albanD</a> , since you’ve helped me before I was wondering if you had any idea? Thank you in advance.</p>

username: albanD
cooked: <p>Hi,</p>
<p>As mentioned  by <a class="mention" href="/u/leviviana">@LeviViana</a>,  <code>self.network.state_dict()[key] = new_state_dict[key]</code> changes the  dictionary created when calling <code>self.network.state_dict()</code> (not the Tensor). So it won’t modify the network itself.<br>
<code>self.network.state_dict()[param_tensor][0][0] = 0</code> does change the original Tensor, because the indexing returns a view of the original Tensor that you modify.<br>
The right way to change the whole Tensor is to change it inplace, not the dict: <code>self.network.state_dict()[key].copy_(new_state_dict[key])</code></p>
<p>Note that (most likely) you don’t want such changes to be tracked by the autograd, so you should wrap these changes (inplace copy) into a <code>with torch.no_grad():</code>.</p>
<p>For your original question, the goal is to load weights that have different sizes from the current ones?<br>
How do you handle that for layers parameters like in_channels/out_channels?</p>

username: Cedric_Gullentops
cooked: <p><a class="mention" href="/u/alband">@albanD</a></p>
<pre><code class="lang-auto">self.network.state_dict()[key].copy_(new_state_dict[key])
</code></pre>
<p>results in a</p>
<pre><code class="lang-auto">RuntimeError: The size of tensor a (32) must match the size of tensor b (13) at non-singleton dimension 0
</code></pre>
<p>since the given tensors have a reduced dimension.</p>
<p>The goal is to load pruned networks (in my case pruned Yolov2 networks) that have been saved during pruning process. The idea was to create a wrapper class that initiates the original network and adjusts it so it can load in the weights of a pruned version of itself, without knowing beforehand what the size of each tensor will be. I hadn’t fully considered the layers parameters yet but I assumed these were accessible from the saved file aswell?</p>

username: albanD
cooked: <p>The saved file only contains the state dict. So only the dictionary of parameters. You won’t find the layers parameters there, unless you save them explicitly on the side.</p>
<p>If you can load the layer informations before and create the PrunedNetwork with the right size, then it might be easier to load it afterwards.</p>

username: Cedric_Gullentops
cooked: <p><a class="mention" href="/u/alband">@albanD</a> Unfortunately this isn’t an option for me anymore since I have a deadline pretty soon and repruning the networks takes several days. I might be able to change each layers parameters based on the size of the tensors within. But I still need a way to edit the original network’s dictionary in this case.</p>

username: albanD
cooked: <p>Editing the dictionnary won’t help. As the dictionnary is created when you call <code>.state_dict()</code>.<br>
If you want to change them, the simplest solution I can think of right now is going to go through the network and for each parameter:</p>
<pre><code class="lang-python"># mod is the current module that has a parameter weight that needs
# to be changed to new_weight
del mod.weight
mod.weight = nn.Parameter(new_weight)

</code></pre>

username: Cedric_Gullentops
cooked: <p><a class="mention" href="/u/alband">@albanD</a>  I managed to to recreate the network and copy over all tensors by running through the keys and assigning them using some string parsing magic to</p>
<pre><code class="lang-auto">self.network.layers[firstindex][secondindex]
</code></pre>
<p>I do however still receive an error:</p>
<pre><code class="lang-auto">RuntimeError: the derivative for 'running_mean' is not implemented
</code></pre>
<p>which I looked up and I understand that these tensors don’t have to be wrapped by nn.Parameter().<br>
However when I leave out the nn.Parameter() I get the following error:</p>
<pre><code class="lang-auto">RuntimeError: Error(s) in loading state_dict for Yolo:
        Unexpected key(s) in state_dict: "layers.0.1_convbatch.layers.1.running_mean", "layers.0.1_convbatch.layers.1.running_var", ...... 
</code></pre>
<p>How do I copy the running var and running mean tensors?</p>

username: albanD
cooked: <p>These are not <code>nn.Parameter</code> but buffers. They are plain Tensors and are registered with <code>register_buffer()</code>.<br>
These buffers should <em>not</em> require gradients. The error you’re seeing is because the formula to compute the gradient wrt the running statistics is not implemented for batchnorm.</p>

username: Cedric_Gullentops
cooked: <p><a class="mention" href="/u/alband">@albanD</a> I don’t think I fully understand. How do I solve this issue?</p>

username: albanD
cooked: <p>To solve “RuntimeError: the derivative for ‘running_mean’ is not implemented”, you need to make sure that <code>batchnorm_mod.running_mean.requires_grad == False</code>.<br>
What might be happening is that in your loading script, you treat both Parameters and buffers the same way, and so you make all of the require gradients. But you want to be careful and make sure that you don’t set <code>requires_grad=True</code> for the buffers.</p>

username: Cedric_Gullentops
cooked: <p><a class="mention" href="/u/alband">@albanD</a></p>
<pre><code class="lang-auto">locallayer.register_buffer('running_mean', new_state_dict[param_tensor])
</code></pre>
<p>This did the trick! Thank you very much again for your help and patience.</p>