Forward_pre_hooks not called after applying nn.uitls.prune methods

ra_schoonhoven: <p>Hello,</p>
<p>I am working with the newly released pruning functionalities in <code>torch.nn.utils.prune</code> and I am working on extending this implementation of the MS-D network:</p>
<aside class="onebox githubfolder" data-onebox-src="https://github.com/ahendriksen/msd_pytorch/tree/master/msd_pytorch">
  <header class="source">

      <a href="https://github.com/ahendriksen/msd_pytorch/tree/master/msd_pytorch" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h3><a href="https://github.com/ahendriksen/msd_pytorch/tree/master/msd_pytorch" target="_blank" rel="noopener nofollow ugc"></a></h3>

  <p><a href="https://github.com/ahendriksen/msd_pytorch/tree/master/msd_pytorch" target="_blank" rel="noopener nofollow ugc">//github.com/ahendriksen/msd_pytorch/tree/master/msd_pytorch</a></p>


  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>This is a network with densely connected 3x3 convolutions followed by a final layer of 1x1 convolutions. Some simplified code:</p>
<pre><code class="lang-auto">import msd_pytorch
import torch.nn.utils.prune as prune

def pytorch_L1_prune(model, depth, perc):
    for k in range(depth):
        wname ="weight"+str(k)
        prune.l1_unstructured(model.msd.msd_block, name=wname, amount=perc)
    prune.l1_unstructured(model.msd.final_layer.linear, name='weight', amount=perc)
    return model

d = 100
model = msd_pytorch.MSDSegmentationModel(c_in, num_labels, d, width, dilations=dilations)
# Load model here
model = pytorch_L1_prune(model, d, 0.3)
</code></pre>
<p><strong>What I have noticed:</strong><br>
The network works fine but I have noticed that the pruning has no effect on the 3x3 convolutions but it does on the 1x1 final layer. My hypothesis is that the <code>forward_pre_hooks</code> are not applied for the <code>forward()</code> pass of the 3x3 convolutions.</p>
<p><strong>Hypothesis:</strong><br>
I found the following note in the docs (<code>torch.nn.modules.module</code>):</p>
<blockquote>
<pre><code>    Although the recipe for forward pass needs to be defined within
    this function, one should call the :class:`Module` instance afterwards
    instead of this since the former takes care of running the
    registered hooks while the latter silently ignores them.
</code></pre>
</blockquote>
<p>I don’t really understand what this note means or how I should resolve it which is why I thought it might be good to clear it up here.</p>
<p>The network has an MSDModule2D which has 2 types of modules; <code>MSDBlock2D</code> for the 3x3 convolutions and <code>MSDFinalLayer</code> at the end and the forward pass is as follows:</p>
<pre><code class="lang-auto">class MSDModule2d(torch.nn.Module):
    def __init__(self, c_in, c_out, depth, width, dilations=[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]):
        super(MSDModule2d, self).__init__()
        ... # initialize dilations etc.
        self.msd_block = MSDBlock2d(self.c_in, self.dilations, self.width)
        self.final_layer = MSDFinalLayer(c_in=c_in + width * depth, c_out=c_out)
        ...

    def forward(self, input):
        output = self.msd_block(input)
        output = self.final_layer(output)
        return output
</code></pre>
<p>The latter basically uses <code>nn.Conv1D()</code> so the above note probably does not apply and it would therefore make sense that <code>forward_pre_hooks</code> are applied.  The <code>MSDBlock2D</code> module implements a somewhat complicated forward pass that I do not want to fully post here:</p>
<pre><code class="lang-auto">class MSDBlockImpl2d(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, dilations, bias, *weights):
          ... # complicated

    @staticmethod
    def backward(ctx, grad_output):
          ... # complicated

class MSDBlock2d(torch.nn.Module):
    def __init__(self, in_channels, dilations, width=1):
        super().__init__()
        ... # initialize weights etc.

    def forward(self, input):
        bias, *weights = self.parameters()
        return MSDBlockImpl2d.apply(input, self.dilations, bias, *weights)
</code></pre>
<p>Any ideas what is going on? Is the above hypothesis correct and if so, what does it mean and what should I add to the <code>MSDBlock2D</code> module?</p>
<p>Kind regards,</p>
<p>Richard</p>
<p><strong>EDIT: CHECK ON PRUNING METHODS</strong><br>
I checked whether the pruning method does its job.</p>
<pre><code class="lang-auto">print(list(model.msd.final_layer.linear.named_buffers()))
print(list(model.msd.msd_block.named_buffers()))
</code></pre>
<p>The model does have the masks etc. Furthermore, the following</p>
<pre><code class="lang-auto">print(model.msd.final_layer.linear._forward_pre_hooks)
print(model.msd.msd_block._forward_pre_hooks)
</code></pre>
<p>shows that the hooks are initialized. They are just not applied in the <code>msd_block</code> case it seems.</p>

albanD: <p>cc <a class="mention" href="/u/michela">@Michela</a> who is the original author of that module.</p>

Michela: <aside class="quote no-group quote-modified" data-username="ra_schoonhoven" data-post="1" data-topic="69497">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/r/e19adc/48.png" class="avatar"> ra_schoonhoven:</div>
<blockquote>
<p>Although the recipe for forward pass needs to be defined within this function, one should call the :class:<code>Module</code> instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.</p>
</blockquote>
</aside>
<p>I’m not the author of that quote, but that states that you should always call <code>module()</code> as opposed to <code>module.forward()</code> if you don’t want to have the hooks ignored. Does the code make an explicit call to <code>forward()</code> by chance? Quickly skimming through it, I found a <code>forward()</code> call here: <a href="https://github.com/ahendriksen/msd_pytorch/blob/master/msd_pytorch/msd_model.py#L213" rel="noopener nofollow ugc">https://github.com/ahendriksen/msd_pytorch/blob/master/msd_pytorch/msd_model.py#L213</a> and here: <a href="https://github.com/ahendriksen/msd_pytorch/blob/master/msd_pytorch/msd_model.py#L250" rel="noopener nofollow ugc">https://github.com/ahendriksen/msd_pytorch/blob/master/msd_pytorch/msd_model.py#L250</a><br>
Not sure how and if these are used at all anywhere in your code. BTW, thanks for finding that quote in the docs and doing lots of research into this yourself before raising this issue.</p>
<p>Another potential problem I see is that in the MSDBlock2d’s forward (<a href="https://github.com/ahendriksen/msd_pytorch/blob/master/msd_pytorch/msd_block.py#L179-L180" rel="noopener nofollow ugc">https://github.com/ahendriksen/msd_pytorch/blob/master/msd_pytorch/msd_block.py#L179-L180</a>) you grab the weights from <code>self.parameters()</code> and then pass them to the complicated <code>forward()</code> above. However, the parameters store the unpruned version of the weight tensors. The pruned versions are just attributes attached to the corresponding Modules. So this might also be the cause of the problem you’re seeing.</p>
<p>Other than these quick observations, from the details you provided (i.e., the fact that the hooks do exist and the net is indeed in a pruned state), it sound as if this is not an issue with pruning, but rather an issue with the specific model you’re using and the way it computes the forward pass, which is probably causing it to ignore the hooks for certain parts of the net. Not sure I’m the best person to debug that but I can take a closer look if needed.</p>

ra_schoonhoven: <p>Hello,</p>
<p>Thanks for the elaborate answer and for taking the time to look into the code. The lines that I had to changes were indeed:</p>
<pre><code class="lang-auto"># In msd_block.py lines 179-180, in the forward pass method
bias, *weights = self.parameters()

# CHANGED TO:
bias = self.bias
weights = (self.__dict__['weight{}'.format(i)] for i in range(len(self.weights)))
</code></pre>
<p>So is it true that when we change the names of the parameters to <code>weights_orig</code> the “pointers” stored in <code>parameters()</code> point to the originals, regardless of the name change? The changed line of code seems to work as pruning is now applied but it does not look very nice. Is this the way it should be done in your opinion?</p>
<p>Kind regards,</p>
<p>Richard</p>
<p>PS: I found the following curious change. The following line runs fine for PyTorch 1.1.0:</p>
<pre><code class="lang-auto"># This runs fine when I run the network (unpruned) in 1.1.0
weights = (self.__getattr__("weight{}".format(i)) for i in range(len(self.weights)))

# But throws the following Error for the pruned network in 1.4.0
AttributeError: 'MSDBlock2d' object has no attribute 'weight0'
</code></pre>
<p>It is probably something minor but I found it curious and therefore changed it to the above. When you print the <code>__dir__()</code> of the object you do see the attributes <code>weight, weight_mask, weight_orig</code> etc.</p>

Michela: <p>Glad we figured out how to make it work!<br>
Yes, the parameters will always be the unpruned versions of the tensors, until you <code>.remove()</code> the reparametrization.</p>
<p>One option: grab both the unpruned parameters from <code>self.parameters()</code> and the masks from <code>self.buffers()</code> and do the multiplication manually (con: would require a new logic block to be executed only if the mask buffers exist).</p>
<p>Another option: without directly accessing <code>__dict__</code>, do something like: <code>weights = [module.weight for module in self.modules() if hasattr(module, 'weight')]</code>, or whatever this needs to be for your model to work. In your specific case, I guess that, since you don’t have single tensors called <code>'weight' </code>in a variety of modules but a single module with a variety of tensors called <code>'weight{i}'</code>, then it would look something like <code>(getattr(self, 'weight{}'.format(i)) for i in range(len(self.weights)))</code>, instead. Should work fine for both pruned and unpruned version of the model.</p>
<p>Re this:</p>
<aside class="quote no-group" data-username="ra_schoonhoven" data-post="4" data-topic="69497">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/r/e19adc/48.png" class="avatar"> ra_schoonhoven:</div>
<blockquote>
<pre><code class="lang-auto"># This runs fine when I run the network (unpruned) in 1.1.0
weights = (self.__getattr__("weight{}".format(i)) for i in range(len(self.weights)))

# But throws the following Error for the pruned network in 1.4.0
AttributeError: 'MSDBlock2d' object has no attribute 'weight0'
</code></pre>
</blockquote>
</aside>
<p>you should do <code>(getattr(self, "weight{}".format(i)) for i in range(len(self.weights)))</code> instead of <code>__getattr__</code>. See this for reference: <a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/module.py#L580" rel="noopener nofollow ugc">https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/module.py#L580</a></p>