Can't reset parameters with torch.nn.utils.prune

username: Ravi_Raju
cooked: <p>Hi, I am trying to implement the lottery ticket hypothesis using torch.nn.utils.prune but I am having trouble trying to reset the model parameters. From what I understand, the pruning module adds a new parameter, <code>weight_orig</code> to each module by overriding <code>weight</code>. When I try to reinitialize the model parameters with <code>model.apply(reinit_weights)</code>, the parameters corresponding to <code>weight_orig</code> in each module is not changing. Below is some code that demonstrates. Does someone have an explanation or potential workaround?</p>
<p>Thanks so much in advance.</p>
<pre><code class="lang-auto"># this allows the model to train with the new optimizer
# but doesn't reset parameters that have the mask added so fails the test_model_change
def reinit_weights(m):
    if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
        m.reset_parameters()

def test_model_change(prev_iter_dict, model):
    for name, param in model.named_parameters():
        prev_param = prev_iter_dict[name]
        assert not torch.allclose(prev_param,param), 'model not updating'

model = LeNet5()
optimizer = optim.Adam(model.parameters(), lr=lr)
model.train().to(device)

pruner = L1Unstructured(0.0)
# adds masks of all ones to each of the layers
for n, m in model.named_modules():
    if isinstance(m, torch.nn.Conv2d):
        pruner.apply(m, name='weight', amount=0.0)
    if isinstance(m, torch.nn.Linear):
        pruner.apply(m, name='weight', amount=0.0)

train(model,1,train_loader,device,optimizer) # train for one epoch to test
full_cap = copy.deepcopy(model.state_dict())
model.apply(reinit_weights) # this reinitalizes the weight parameter but doesn't change the new_weight parameter
test_model_change(full_cap, model) # passes for conv0.bias but fails for conv0.weight
</code></pre>

username: alx
cooked: <p>Personally, I simply save the weights I’d like to re-initialize to and then load the .pt file when its time to re-initialize.</p>

username: Ravi_Raju
cooked: <p>Thank you for your response but this doesn’t work for me. This passes my test_model_change function so the weights do get reinitialized but when I try to retrain my model, the loss does not decrease. Interestingly, when I reset the model to the original weight parameters at the first initialization I don’t face this issue, only when I randomly try to set weights to some values.</p>
<pre><code class="lang-auto">def reset_weights_rand(model):
    rand_state_dict = pickle.load(open('rand_dict.pkl', 'rb'))
    for name, param in model.named_parameters():
        param.data = rand_state_dict[name]
rand_state_dict = copy.deepcopy(initial_state_dict)

for key,value in rand_state_dict.items():
    if 'bias' in key:
        init.normal_(value.data)
    else:
        init.xavier_normal_(value.data)        

with open(r"rand_dict.pkl", "wb") as output_file:
    pickle.dump(rand_state_dict, output_file)

train(model,1,train_loader,device,optimizer) # train for one epoch to test

full_cap = copy.deepcopy(model.state_dict())

reset_weights_rand(model)

# model.apply(reinit_weights) # this reinitalizes the weight parameter but doesn't change the new_weight parameter
test_model_change(full_cap, model) # passes now

# train 90% capacity for 10 more epochs
optimizer2 = torch.optim.Adam(model.parameters(), lr=lr) # problematic

train(model,5, train_loader,device,optimizer2)
</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/1/b164848d36b5edf7903adf4b84b68bea5ba12e02.png" data-download-href="https://discuss.pytorch.org/uploads/default/b164848d36b5edf7903adf4b84b68bea5ba12e02" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/b/1/b164848d36b5edf7903adf4b84b68bea5ba12e02_2_516x500.png" alt="image" data-base62-sha1="pjhRYEYpQol9BPpPe1Hs27hOhHQ" width="516" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/b/1/b164848d36b5edf7903adf4b84b68bea5ba12e02_2_516x500.png, https://discuss.pytorch.org/uploads/default/optimized/3X/b/1/b164848d36b5edf7903adf4b84b68bea5ba12e02_2_774x750.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/b/1/b164848d36b5edf7903adf4b84b68bea5ba12e02_2_1032x1000.png 2x" data-dominant-color="EDEEEE"><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">image</span><span class="informations">1530×1480 342 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><br>
As you can see the loss does not decrease even though no pruning as been done.</p>

username: Ravi_Raju
cooked: <p>Hi,</p>
<p>I actually managed to hack a workaround with the tools in <code>torch.nn.utils.prune</code>. The trick is save the dictionary of masks and remove the pruning modules with <code>remove</code> and then add them back with <code>CustomFromMask</code>.</p>
<pre><code class="lang-auto">from torch.nn.utils.prune import l1_unstructured, remove, CustomFromMask,is_pruned


def add_masks(model,masks):
    mask_pruner = CustomFromMask(None)
    for module_name, module in model.named_modules():
        key = f"{module_name}.weight_mask"
        if key in masks:
            if isinstance(module, torch.nn.Conv2d):
                _mask = masks[key]
                mask_pruner.apply(module, 'weight', _mask)
            if isinstance(module, torch.nn.Linear):
                _mask = masks[key]
                mask_pruner.apply(module, 'weight', _mask)

def merge_masks(model):
    for n, m in model.named_modules():
        if is_pruned(m)==False:
            continue
        if isinstance(m, torch.nn.Conv2d):
            remove(m, name='weight')
        if isinstance(m, torch.nn.Linear):
            remove(m, name='weight')
</code></pre>
<p>With this fix, I can reinitialize my network with no issues. Hopefully this can help someone if they were dealing with my issue.</p>