Custom pruning method for Global Unstructured Pruning

Shashank_Priyadarshi: <p>Hi all,<br>
So I am trying to add a custom global pruning method by extending PyTorch pruning class. So I already have the mask for each layer stored in the dict where key is the module name like conv1 and the value is the masks for the specific module. But I want to use the torch class so that in future when I sometime the model then also those weights remain 0.<br>
Any idea how can I do it. It will be a great help.</p>

alicerichard65: <blockquote>
<p>Hi all,<br>
So I am trying to add a custom global pruning method by extending PyTorch pruning class. So I already have the mask for each layer stored in the dict where key is the module name like conv1 and the value is the masks for the specific module. But I want to use the torch class so that in future when I sometime the model then also those weights remain 0 <a href="https://www.litebluelite.com" rel="noopener nofollow ugc">learn more</a><br>
Any idea how can I do it. It will be a great help.</p>
</blockquote>
<p>To achieve this, you can extend the PyTorch pruning class and use the <code>torch.nn.utils.prune</code> module. Here’s a high-level approach:</p>
<ol>
<li><strong>Create a Custom Pruning Class:</strong></li>
</ol>
<ul>
<li>Extend the <code>PruningMethod</code> class from <code>torch.nn.utils.prune</code>.</li>
<li>Implement your custom pruning logic in the <code>compute_mask</code> method.</li>
</ul>
<ol start="2">
<li><strong>Apply the Custom Pruning Method:</strong></li>
</ol>
<ul>
<li>Use the <code>prune</code> module to apply your custom pruning method to the model parameters.</li>
</ul>
<p>Here’s a basic example to get you started:</p>
<pre><code class="lang-auto">import torch
import torch.nn.utils.prune as prune

class CustomPruning(prune.PruningMethod):
    def compute_mask(self, t, default_mask):
        # Your custom logic to compute the mask
        mask = torch.zeros_like(t)
        # Example: Set mask to 0 for weights below a certain threshold
        mask[t.abs() &lt; 0.1] = 1
        return mask

# Example model
model = torch.nn.Sequential(
    torch.nn.Linear(10, 5),
    torch.nn.ReLU(),
    torch.nn.Linear(5, 1)
)

# Apply custom pruning
parameters_to_prune = (model[0], 'weight')
prune.global_unstructured(parameters_to_prune, pruning_method=CustomPruning(), amount=0.2)

</code></pre>
<p>This example demonstrates how to create a custom pruning class and apply it to a model. You can adjust the <code>compute_mask</code> method to use your stored masks and ensure the weights remain zero when reloading the model.</p>

Shashank_Priyadarshi: <p>Hey thanks, I done something like this - I have one more question:</p>
<p>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>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.</li>
<li>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</li>
</ol>