Pruning the Convnets

gsp-27: <p>Hello,</p>
<p>I am trying to implement weight pruning using forward hook, but somehow it is giving me invalid syntax, below is the code</p>
<pre><code class="lang-python">import torch
import torch.nn as nn
from torch.autograd import Variable
import numpy as np

# this code implements pruning using register buffer feature to save input mask

def compute_mask(weights):
    thresh = weights.std()
    m1 = weights &gt; thresh
    m2 = weights &lt; (-thresh)
    mask = torch.ones(weights.size())
    mask = mask-m1.float()
    mask = mask-m2.float()
    return mask

class PrunedSqueezenet(nn.Module):
    def __init__(self, to_prune, pretrained_weight):
        """
        takes a list of layers to prune, model, weights
        to_prune: a list of all the layers on which pruning should be applied
        model: architecture of the model
        weights: pretrained weights to use for the model
        """
        super(self, PrunedSqueezenet).__init__()
        self.to_prune = to_prune
        # get the model ready
        self.base_model = model.SqueezeNet()
        pretrained_weights = torch.load(pretrained_weight)
        base_model.load_state_dict(pretrained_weights)

        self.layers = self.base_model._modules.keys()
        # compute the mask for the weights
        for l in to_prune:
            if "fire" in l:
                curr_layer = self.base_model._modules.get(l)._modules.get('conv3')
                weights = curr_layer.weight.data
                # save the mask
                curr_layer.register_buffer('mask', compute_mask(weights))
                # change the computed output of conv3 layer in the fire
                curr_layer.register_forward_hook(
                    lambda m, i, o: \
                        print("Hello this is ok")
                )

            elif "conv" in l:
                curr_layer = self.base_model._modules.get(l)
                weights = curr_layer.weight.data
                # save the mask
                curr_layer.register_buffer('mask', compute_mask(weights))
                # change the computed output of conv3 layer in the fire
                curr_layer.register_forward_hook(
                    lambda m, i, o: \
                        print("Hello this is ok"))
                )
            
            else:
                print("I dont understand what you are talking about")

    def forward(self, x):
        return self.base_model(x)

if __name__ == '__main__':
    net = PrunedSqueezenet(to_prune=['fire9'], pretrained_weight='pretrained_models/squeezedet_compatible.pth')
    x = Variable(torch.randn(1, 3, 32, 32))
    print(net(x))
</code></pre>

smth: <p>what’s the exact error and stack trace?</p>
<p>invalid syntax probably means invalid python syntax.</p>

gsp-27: <p>I solved it, it was some python syntax error. Sorry for the trouble.</p>
<p>But I do have a logical question in “pruning”. So my line of thinking is I will compute mask for the specified layers and store it as the layer buffer and then implement a hook which will multiply the weight matrix with the mask thus essentially making the weights zero at that position.</p>
<p>And hence they will not contribute to the output… The question is whether this hook is applied after the forward of the layer or before, because in my before will make sense.</p>
<p>If this thinking is not right direction, would you please suggest me a way to achieve the desired behaviour?</p>

smth: <p>The hook is applied after the forward of the layer.</p>

apaszke: <p>I think it would be simpler to set the pruned weights to 0 and mask gradients after backward(), so that they’re never modified.</p>

gsp-27: <p>Thanks I shall try the method of <a class="mention" href="/u/apaszke">@apaszke</a></p>

angup143: <p>I’m in a similar situation where I need to mask weights during forward(). I cannot set them to zero initially because the weight_mask gets updated every iteration based on the values of the weights and the gradients during backprop update all the weights (not just the unpruned ones).</p>
<p>Any suggestions on how I could do that short of creating new layers?</p>

evcu: <p>I would suggest using a wrapper module like the one here:<br>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/evcu/pytorchpruner/blob/76b3fb7ded9af7c6a84d4792efdaed235cb0dd06/pytorchpruner/modules.py#L4" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/evcu/pytorchpruner/blob/76b3fb7ded9af7c6a84d4792efdaed235cb0dd06/pytorchpruner/modules.py#L4" target="_blank" rel="nofollow noopener">evcu/pytorchpruner/blob/76b3fb7ded9af7c6a84d4792efdaed235cb0dd06/pytorchpruner/modules.py#L4</a></h4>
<pre class="onebox"><code class="lang-py"><ol class="start lines" start="1" style="counter-reset: li-counter 0 ;">
<li>import torch</li>
<li>from torch.nn.modules import Module</li>
<li>
</li>
<li class="selected">class MaskedModule(Module):</li>
<li>r"""Implements masked module for prunning etc...</li>
<li>it creates a mask for each layer and holds those masks inside a dictionary belongs to</li>
<li>this module</li>
<li>"""</li>
<li>DEFAULT_MASKED_MODULES = (torch.nn.Conv2d,torch.nn.Linear)</li>
<li>def __init__(self, module, masked_modules=DEFAULT_MASKED_MODULES):</li>
<li>    super(MaskedModule, self).__init__()</li>
<li>    self.module = module</li>
<li>    self._mask_dict = {} # nn.Module-&gt;torch.ByteTensor</li>
<li>    self._inp_dict = {} # nn.Module-&gt;torch.Tensor</li>
</ol></code></pre>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>
</p>
<p>The library itself is an ongoing project of mine about pruning and calculating saliencies.</p>

wxwx: <p>hello, i am interesting in your code. I have a question : how can i retrain a neural network without change some fixed weight? Hope for you reply! thanks!!!</p>

evcu: <p>One method is to use a method to mask the gradients after <code>backward()</code> call and before <code>step()</code>. I use that one in my code.</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/evcu/pytorchpruner/blob/76b3fb7ded9af7c6a84d4792efdaed235cb0dd06/pytorchpruner/modules.py#L29" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/evcu/pytorchpruner/blob/76b3fb7ded9af7c6a84d4792efdaed235cb0dd06/pytorchpruner/modules.py#L29" target="_blank" rel="nofollow noopener">evcu/pytorchpruner/blob/76b3fb7ded9af7c6a84d4792efdaed235cb0dd06/pytorchpruner/modules.py#L29</a></h4>
<pre class="onebox"><code class="lang-py"><ol class="start lines" start="19" style="counter-reset: li-counter 18 ;">
<li>
</li>
<li>def initialize_transparent_masks(self):</li>
<li>    def helper(m):</li>
<li>        if isinstance(m, self.DEFAULT_MASKED_MODULES):</li>
<li>            self._mask_dict[m] = torch.ones(m.weight.data.size())</li>
<li>    self.module.apply(helper)</li>
<li>
</li>
<li>def forward(self, *inputs, **kwargs):</li>
<li>    return self.module.forward(*inputs, **kwargs)</li>
<li>
</li>
<li class="selected">def apply_mask_on_gradients(self):</li>
<li>    def helper(m):</li>
<li>        if isinstance(m, self.DEFAULT_MASKED_MODULES):</li>
<li>            m.weight.grad.data[self._mask_dict[m]!=1]=0</li>
<li>    self.module.apply(helper)</li>
<li>
</li>
<li>def initiliaze_forward_hooks(self):</li>
<li>    def f_hook(module,inp,out):</li>
<li>        #TODO check inp is a tuple of size 1.</li>
<li>        self._inp_dict[module]=inp[0]</li>
<li>    def helper(m):</li>
</ol></code></pre>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<p>A second method could be doing the same with backwardhooks; where you use <code>module.register_backward_hook()</code> to mask the gradients. You can have a flag in your wrapper module like <code>is_masked=True/False</code> and then you can check that flag first thing in your hook: <code>if module.is_masked: do masking else: pass</code></p>
<p>One tricky thing is to reinitialize your optimizer or gradient history of it after each update on your mask to prevent momentum to come into play.</p>

wxwx: <p>wow,very thanks for reply!!! i will try . thanks</p>

wxwx: <p>hello, i am sorry to disturb you again. i have the same question:<br>
I set a masks=<span class="chcklst-box fa fa-square-o fa-fw"></span> to save which weight equal to zero and don’t change these weights.<br>
<img src="https://discuss.pytorch.org/uploads/default/original/2X/d/dc4966b763828623821893c4e81fc038e18233e8.png" alt="mask" data-base62-sha1="vqKiU7fWsP5NXGMYjIHfa3fhBu8" width="650" height="133"><br>
second, i set them to variable which needn’t grad<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/2X/7/7995030b2fbeb03b1505bb01e88cac4ab9044944.png" data-download-href="https://discuss.pytorch.org/uploads/default/7995030b2fbeb03b1505bb01e88cac4ab9044944" title="mask2"><img src="https://discuss.pytorch.org/uploads/default/original/2X/7/7995030b2fbeb03b1505bb01e88cac4ab9044944.png" alt="mask2" data-base62-sha1="hlz1SX94f3k8sZfaEVbzJh6s04I" width="690" height="45" data-dominant-color="0E1111"><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">mask2</span><span class="informations">926×61 2.69 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>
my question is how to apply masks to prevent autograd change these weights</p>

evcu: <p>Efficient way of doing this is having byte_tensor mask with same size as the original <code>.weight</code> tensor. So you can just remove <code>.float()</code> casting. Your mask would have ones for pruned weights. You don’t need a Variable of masks. You can just use the masks for indexing like below for example.</p>
<pre><code class="lang-auto">loss.backward()
for layer,mask in zip(X,Y):
     layer.weight.grad.data[mask]=0
optimizer.step()
</code></pre>
<p>You just need to keep track which mask belongs to which weight. That’s MaskedModule wrapper is doing in the code I shared above.</p>

wxwx: <p>very thanks for reply, i will try again ! thanks</p>

wxwx: <p>I have a question about your code. whether the <strong>loss.backward()</strong> calculate zeroed weight? my purpose is when call the <strong>loss.backward()</strong>, zeroed weight don’t participate in backpropagation. sorry to disturb you again:blush:</p>

wxwx: <p>I have a question about your code. whether the <strong>loss.backward()</strong> calculate zeroed weight? my purpose is when call the <strong>loss.backward()</strong>, zeroed weight don’t participate in backpropagation. sorry to disturb you again</p>

Daniel_Ortega: <p>Hi, do you find a solution for this?</p>