Reset model weights

username: unnir
cooked: <p>I would like to know, if there is a way to reset weights for a PyTorch model.</p>
<p>Here is my code:</p>
<pre><code class="lang-auto">class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 16, kernel_size=5)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=5)
        self.conv3 = nn.Conv2d(32, 64, kernel_size=5)
        self.conv4 = nn.Conv2d(64, 64, kernel_size=5)

        self.pool = nn.MaxPool2d(5, stride=3)
        self.pool2 = nn.MaxPool2d(3, stride=1)
        self.activation = nn.ReLU()
        self.fc1 = nn.Linear(4096, 20)
        self.fc2 = nn.Linear(20, 1)

    def forward(self, x):
        x = self.activation( self.pool(self.conv1(x)) )
        x = self.activation( self.pool(self.conv2(x)) )
        x = self.activation( self.pool(self.conv3(x)) )
        x = self.activation( self.pool2(self.conv4(x)) )
        x = x.view(-1, 64*8*8)

        x = self.activation(self.fc1(x))
        x = self.fc2(x)
        return x
</code></pre>
<p>I just to want average couple runs of the model, in order to evaluate it.</p>
<p>Any idea how can I do that?</p>

username: ptrblck
cooked: <p>You could save the <code>state_dict</code> and load it for resetting the model. Have a look at the <a href="https://pytorch.org/docs/master/notes/serialization.html">Serialization Semantics</a> to see how to do it.<br>
Would this work for you or do you want to re-initialize it to random weights?</p>

username: unnir
cooked: <p>thank you for the hint.</p>
<p>However, is there a way for the random re-initialization?</p>

username: ptrblck
cooked: <p>Sure! You just have to define your init function:</p>
<pre><code class="lang-auto">def weights_init(m):
    if isinstance(m, nn.Conv2d):
        torch.nn.init.xavier_uniform(m.weight.data)
</code></pre>
<p>And call it on the model with:</p>
<pre><code class="lang-auto">model.apply(weight_init)
</code></pre>
<p>If you want to have the same random weights for each initialization, you would need to set the seed before calling this method with:</p>
<pre><code class="lang-auto">torch.manual_seed(your_seed)
</code></pre>

username: Brando_Miranda
cooked: <p>This probably works too:</p>
<hr>
<p>This is not a robust solution and wont work for anything except core torch.nn layers, but this works:</p>
<pre><code class="lang-auto">for layer in model.children():
   if hasattr(layer, 'reset_parameters'):
       layer.reset_parameters()
</code></pre>
<hr>
<p>credit: <a href="https://discuss.pytorch.org/t/reinitializing-the-weights-after-each-cross-validation-fold/11034/2" class="inline-onebox">Reinitializing the weights after each cross validation fold</a></p>

username: Frida_de_Sigley
cooked: <p>could I perhaps double check if below is a robust solution? it is working - However I does seem useful to corroborate. I am looking for a solution that clears all weights between iterations of a hyper-parameter search. I am running individual models as sub-process’.</p>
<pre><code class="lang-auto">chk_dir = '/root/.cache/torch/hub/checkpoints/'

if os.path.isdir(chk_dir):
    for chkpnt in os.scandir(chk_dir):
        print(f'rm"ing {chkpnt.path}')
        os.system(f'rm {chkpnt.path}')
</code></pre>

username: Brando_Miranda
cooked: <p>why doesn’t this work for you:</p>
<aside class="quote" data-post="1" data-topic="20819">
  <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/sohrab_salimian/48/2326_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/how-to-re-set-alll-parameters-in-a-network/20819">How to re-set alll parameters in a network</a> 
  </div>
  <blockquote>
    How to re-set the weights for the entire network, using the original pytorch weight initialization
  </blockquote>
</aside>

<p><a class="mention" href="/u/unnir">@unnir</a></p>
<p>Here is the code with an example that runs:</p>
<pre><code class="lang-auto">def lp_norm(mdl: nn.Module, p: int = 2) -&gt; Tensor:
    lp_norms = [w.norm(p) for name, w in mdl.named_parameters()]
    return sum(lp_norms)

def reset_all_weights(model: nn.Module) -&gt; None:
    """
    refs:
        - https://discuss.pytorch.org/t/how-to-re-set-alll-parameters-in-a-network/20819/6
        - https://stackoverflow.com/questions/63627997/reset-parameters-of-a-neural-network-in-pytorch
        - https://pytorch.org/docs/stable/generated/torch.nn.Module.html
    """

    @torch.no_grad()
    def weight_reset(m: nn.Module):
        # - check if the current module has reset_parameters &amp; if it's callabed called it on m
        reset_parameters = getattr(m, "reset_parameters", None)
        if callable(reset_parameters):
            m.reset_parameters()

    # Applies fn recursively to every submodule see: https://pytorch.org/docs/stable/generated/torch.nn.Module.html
    model.apply(fn=weight_reset)


def reset_all_linear_layer_weights(model: nn.Module) -&gt; nn.Module:
    """
    Resets all weights recursively for linear layers.

    ref:
        - https://pytorch.org/docs/stable/generated/torch.nn.Module.html
    """

    @torch.no_grad()
    def init_weights(m):
        if type(m) == nn.Linear:
            m.weight.fill_(1.0)

    # Applies fn recursively to every submodule see: https://pytorch.org/docs/stable/generated/torch.nn.Module.html
    model.apply(init_weights)


def reset_all_weights_with_specific_layer_type(model: nn.Module, modules_type2reset) -&gt; nn.Module:
    """
    Resets all weights recursively for linear layers.

    ref:
        - https://pytorch.org/docs/stable/generated/torch.nn.Module.html
    """

    @torch.no_grad()
    def init_weights(m):
        if type(m) == modules_type2reset:
            # if type(m) == torch.nn.BatchNorm2d:
            #     m.weight.fill_(1.0)
            m.reset_parameters()

    # Applies fn recursively to every submodule see: https://pytorch.org/docs/stable/generated/torch.nn.Module.html
    model.apply(init_weights)


# -- tests

def reset_params_test():
    import torchvision.models as models
    from uutils.torch_uu import lp_norm

    resnet18 = models.resnet18(pretrained=True)
    resnet18_random = models.resnet18(pretrained=False)

    print(f'{lp_norm(resnet18)=}')
    print(f'{lp_norm(resnet18_random)=}')
    print(f'{lp_norm(resnet18)=}')
    reset_all_weights(resnet18)
    print(f'{lp_norm(resnet18)=}')


if __name__ == '__main__':
    reset_params_test()
    print('Done! \a\n')
</code></pre>
<p>output:</p>
<pre><code class="lang-auto">lp_norm(resnet18)=tensor(517.5472, grad_fn=&lt;AddBackward0&gt;)
lp_norm(resnet18_random)=tensor(668.3687, grad_fn=&lt;AddBackward0&gt;)
lp_norm(resnet18)=tensor(517.5472, grad_fn=&lt;AddBackward0&gt;)
lp_norm(resnet18)=tensor(476.0836, grad_fn=&lt;AddBackward0&gt;)
Done!
</code></pre>
<p>I am assuming this works because I calculated the norm twice for the pre-trained net and it was the same both times before calling reset.</p>
<p>Though I was unhappy it wasn’t closer to the norm of the random net I must admit but I think this is good enough.</p>
<p>related: <a href="https://stackoverflow.com/a/69905671/1601580" class="inline-onebox" rel="noopener nofollow ugc">python 3.x - Reset parameters of a neural network in pytorch - Stack Overflow</a></p>

username: Sourabh
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> <a class="mention" href="/u/brando_miranda">@Brando_Miranda</a>  I was trying to do something with regards to resetting the weights and then applying a new tensor as my weights for a particular layer. I just had one doubt regarding the above discussion -  does the <code>reset_parameters()</code> also clear all the associated memory that the layer occupies?</p>

username: ptrblck
cooked: <p><code>reset_parameters</code> usually changes the parameters inplace as seen for e.g. <a href="https://github.com/pytorch/pytorch/blob/702d375df524565f8f2741c8d96df254e450653c/torch/nn/modules/linear.py#L92-L100"><code>nn.Linear</code></a> so the memory usage shouldn’t change.</p>

username: oustella
cooked: <p>How is <code>reset_parameters</code> different from re-instantiate the model instance?</p>

username: ptrblck
cooked: <p><code>.reset_parameters()</code> will reset the parameters inplace, such that the actual parameters are the same objects but their values will be manipulated.<br>
This would allow you to use the same optimizer etc. in case you’ve already passed the parameters to it.<br>
If you are creating a new module, you would of course also reset the parameters, but these parameters are new objects which you might need to pass to an optimizer again (depending on your actual use case).</p>

username: nikita_gordia
cooked: <p>Actually there is more flexible variant:</p>
<pre><code class="lang-auto">def reset_model_weights(layer):
    if hasattr(layer, 'reset_parameters'):
        layer.reset_parameters()
    else:
        if hasattr(layer, 'children'):
            for child in layer.children():
                reset_model_weights(child)
</code></pre>

username: oguz-hanoglu
cooked: <p>typo: weight_init → weights_init</p>

username: Filippa
cooked: <p>Is it possible to apply Xavier (or a similar initialization) to only a subset of weights in a layer? For example, I want to randomly re-initialize 20% of the weights while keeping the rest unchanged. Is there a standard way to do this in PyTorch, or do you have any suggestions?</p>

username: J_Johnson
cooked: <p><a class="mention" href="/u/filippa">@Filippa</a></p>
<p>You can apply a random boolean mask on any tensor in PyTorch, which will then only make the changes where <code>True</code> exists.</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.nn.init as init

def partial_xavier_init(model: nn.Module, percent: float):
    """
    Applies Xavier uniform initialization to a random subset (percent) of weights in each layer.
    
    Args:
    - model: The PyTorch model.
    - percent: Fraction of weights to reinitialize (0.0 to 1.0).
    """
    for name, param in model.named_parameters():
        if 'weight' in name:  # Apply only to weight parameters
            # Create a tensor with Xavier initialization
            new_param = torch.empty_like(param)
            init.xavier_uniform_(new_param)
            
            # Create random mask
            mask = torch.rand_like(param) &lt; percent
            
            # Apply masked update
            param.data[mask] = new_param[mask]
</code></pre>