Finetuning a model after pruning - Autograd question

mzimmer: <p>I want to prune after training and then finetune the pruned model. If I use the <code>torch.nn.utils.prune</code> library, as far as I understand it, during the forward pass the weights of a layer will first be zeroed using the pruning mask (via pre forward hook). This however makes the masking part of the backward step and it will have an effect on the actual gradient updates.</p>
<p><strong>What I want to do is the following:</strong><br>
I want to prune a model and then continue training by just ignoring the pruned weights, similar to as they would have been removed. The mask works correctly in the forward pass as all pruned weights are set to 0, but wouldn’t I get different gradients when doing backpropagation? How can I do this within the pytorch library?</p>
<p>Thanks a lot!</p>

Song_zy: <p>Hi Max,<br>
Have you solved this problem? I got the same problem even I trying to mask those pruned weight in every iteration. I hope we can discuss on this problem if you have any idea.</p>
<p>Thank you!</p>

mzimmer: <p>Hey, yes actually the prune module behaves just as I wanted it to behave. If you prune a module, than the weights will be zero in the forward pass and the gradients will be zero as well.</p>
<p>Does that answer your question?<br>
Best,<br>
Max</p>

davidweb: <p>Hey Max,<br>
so reading your reply, I understand that if we do not want to update pruned weights during fine-tuning the pruning module takes care of that?<br>
So gradients of pruned weights are zero by default?</p>

mzimmer: <p>Yes, exactly. Gradients of pruned weights will be zero (I experimentally verified this).<br>
Best,<br>
Max</p>

davidweb: <p>Hey <a class="mention" href="/u/mzimmer">@mzimmer</a> ,<br>
again I would appreciate your help as I am trying to prune iteratively right now:<br>
From my Conv modules, I pruned the weights using <code>nn.utils.prune</code>.<br>
Now I want to finetune the model again - more specifically only the Conv modules, as the Linear modules were not pruned. Which parameters do I pass to the optimizer? <code>weight_orig</code> or <code>weights</code>?</p>
<p>Because <code>model.named_parameters()</code> returns only <code>weight_orig</code> and <code>bias</code> of each module.<br>
The pruned <code>weight</code> is now accessible directly by calling <code>module.weight</code>.</p>
<p>So which parameters do I need to pass to the optimizer when I want to finetune the model again? Because I want pruned weights (which are basically just weights with the value zero) to stay zero.<br>
My intuition says the pruned <code>weights</code> need to be passed now - how can I do that easily?<br>
Maybe removing the <code>weight_orig</code> of the module using the <code>remove</code> function of the <code>nn.utils.prune</code> library?</p>
<p>Happy to hear your thoughts - maybe you can help me out.<br>
Thanks!</p>
<p>UPDATE: I tried removing the <code>weight_orig</code> of the module using the <code>remove</code> function of the <code>nn.utils.prune</code>. Works in theory, but weights that were pruned (= set to zero) were updated during training again - resulting in non-zero weights again. So this cannot be the solution.</p>

mzimmer: <p>Hey <a class="mention" href="/u/davidweb">@davidweb</a>,<br>
so I am not a 100% sure on this, but as soon as you call <code>remove()</code>, the zero weights will get updates again.</p>
<p>Have you tried the following:</p>
<ol>
<li>You should just pass <code>model.parameters()</code> to the optimizer, as usual.</li>
<li>If I understand you correctly, you don’t want the Linear modules to be updated during finetuning, right? Just continue training the full model, but freeze the layers that you don’t want to be updated, check for example <a href="https://discuss.pytorch.org/t/how-the-pytorch-freeze-network-in-some-layers-only-the-rest-of-the-training/7088">this link</a>.</li>
</ol>
<p>I think this should do what you want. As long as you do not remove the pruning, pruned weights won’t get updated. Freezing allows you to disable certain layers of your network. Let me know whether this fixes your issue.</p>

davidweb: <p>Hey <a class="mention" href="/u/mzimmer">@mzimmer</a>,</p>
<p>yes exactly, freezing FC Layers during finetuning of a pruned model is what I wanted.<br>
Using your suggestion worked perfectly - after several iterations zero weights stay zero.</p>
<p>Thanks for helping me out! Very much appreciated.</p>

Hajar_Mazaheri: <p>Hi Max,<a class="mention" href="/u/mzimmer">@mzimmer</a><br>
Is the size of the pruned model (due to weight-mask storage) larger than the original model? I pruned the encoder layers of my model and after saving I noticed that the model size got bigger.</p>
<p>‌Best<br>
Thank you</p>

mzimmer: <p>Hi <a class="mention" href="/u/hajar_mazaheri">@Hajar_Mazaheri</a> ,<br>
yes, this will increase the storage requirements, since for every pruned tensor you will need to store the mask as well.</p>

Hajar_Mazaheri: <p>Thank you for your response, <a class="mention" href="/u/mzimmer">@mzimmer</a><br>
In your opinion, is pruning the weights one of the compression techniques, or is the model lightweight? And do pruning the weights reduce computational complexity(MAC) during training and inference?<br>
I observed that after pruning the weights and retraining the model, the results improved.</p>
<p>I would be happy to know your opinion and experience on this matter.</p>

mzimmer: <p>I am not sure what you mean by pruning being one of the compression techniques or the model being lightweight. Could you elaborate?</p>
<p>Whether pruning reduces the number of FLOPs or MACs depends heavily on a number of things, e.g., when it is applied (post-training or during/before training), whether the sparsity is structured (removing entire filters or just weights without a clear pattern in the tensor), the device used (gpu or cpu), and many others. I recommend reading into the survey by Hoefler et al.: <a href="https://arxiv.org/abs/2102.00554" class="inline-onebox" rel="noopener nofollow ugc">[2102.00554] Sparsity in Deep Learning: Pruning and growth for efficient inference and training in neural networks</a></p>
<p>In my experience, retraining can drastically improve the performance of the model. If this interests you, one of my papers on retraining might clarify some terminology in this regard: <a href="https://arxiv.org/abs/2111.00843" class="inline-onebox" rel="noopener nofollow ugc">[2111.00843] How I Learned to Stop Worrying and Love Retraining</a></p>
<p>Best,<br>
Max</p>

Hajar_Mazaheri: <p>By compression methods, I mean methods that can help to have a lighter model for inference applications, especially for edge devices. I explored methods such as quantization and knowledge distillation. And now I am investigating the pruning method for my model because in terms of optimization methods, pruning is one of the most effective techniques.<br>
And I want to know that zeroing the weights in addition to improving the results can increase the speed of inference (or training) in the model considering that the size of the model is also increased.</p>
<p>Thank you for your guidance<br>
Your articles can be a good guide for me in this field</p>

mzimmer: <p>Pruning can make your model more efficient, but as said, it also depends on the type of pruning. If you prune in an unstructured way (as magnitude pruning would do), then this sparsity is in general not directly exploitable. More common approaches that result in direct speedups are structured pruning strategies, such as convolutional filter pruning, pruning of entire neurons, or semi-structured sparsities like 2:4 or 4:8 sparsities. Hope that this answers the question, otherwise feel free to check out the mentioned articles for more details!</p>

Hajar_Mazaheri: <p>Yes, exactly. <a class="mention" href="/u/mzimmer">@mzimmer</a><br>
Your detailed explanation clarified the nuances of pruning techniques, highlighting the importance of structured approaches for achieving direct speedups. I appreciate your insights into the distinctions between unstructured and structured pruning methods, such as convolutional filter pruning and sparsities like 2:4 or 4:8.</p>
<p>Thank you for your expertise and assistance.</p>

Hajar_Mazaheri: <p>Hi Max <a class="mention" href="/u/mzimmer">@mzimmer</a><br>
Your suggested articles give me a lot of information about pruning. Considering that you have a lot of information in this field, I have one more question for you.<br>
When we use unstructured (as magnitude pruning) pruning, the zeroed weights are updated during training, and after retraining(fine-tuning), the zeroed weights have a new value, and our model structure is not a sparse structure. Is it true?<br>
Do you think this pruning can still improve model performance if the zeroed weights are frozen during training?</p>
<p>Thank you!</p>

Hajar_Mazaheri: <p>I found the answer to my question. It is related to differentiate between static and dynamic sparsity during training.<br>
Dynamic sparsity during training starts with schemes that iteratively prune and add (regrow) elements during the training phase. but  Fixed sparsity during training, this structure can be hand-tuned such as “structured sparsity” for transformers. <img src="https://discuss.pytorch.org/images/emoji/apple/blush.png?v=12" title=":blush:" class="emoji" alt=":blush:" loading="lazy" width="20" height="20"></p>

Mayank_Kumar: <p><a class="mention" href="/u/hajar_mazaheri">@Hajar_Mazaheri</a> have you came across papers that do both these kinds of pruning ? Any ones with open source code ?</p>

Shashank_Priyadarshi: <p><a class="mention" href="/u/mzimmer">@mzimmer</a><br>
I have one question, I pruned the model that has now weight_masks and weight_orig and an attribute weight. So during finetuning I want to keep the pruning such that only those weights are updated which are not 0. So I found 2 ways of doing it, correct me if I am wrong -</p>
<ol>
<li>
<p>Finetune the pruned model directly, here the question is are weight_orig get updated or weight(attribute)? And if weight_orig updates then the final weight(parameter) is something like weight=weight_orig*weight_mask. And when I save this model then I have weight_orig and weight_mask in the state dict, so once I want to use this fine-tuned pruned model again then in that case I will load the model and do .weight as weight_orig*weight_mask.</p>
</li>
<li>
<p>After fine-tuning should I apply Reparametrization? so when I am loading the saved model again then I only get pruned weights in my .weight</p>
</li>
</ol>