Understanding `register_forward_pre_hook` and `register_backward_hook`

username: MrRobot
cooked: <h2>Problem</h2>
<p>I wrote following code to check my understanding about these two functions</p>
<pre><code class="lang-python">class LinearTransformation(nn.Module):
    def __init__(self, constant=False):
        super(LinearTransformation, self).__init__()
        self.linear = nn.Linear(2, 2)
        if constant:
            nn.init.constant(self.linear.weight, 10)
        else:
            nn.init.eye_(self.linear.weight)
        nn.init.constant_(self.linear.bias, 0)
    
    def forward(self, x):
        return torch.norm(self.linear(x)) ** 2
</code></pre>
<p>Invoking this simple class in the following snippets return</p>
<pre><code class="lang-python">def hook(self, input):
    print(input)
model = LinearTransformation(constant=False)

x = torch.tensor([1., 2.])

model.linear.register_forward_pre_hook(hook)
model(x)
# (tensor([1., 2.]),)
# tensor(5., grad_fn=&lt;PowBackward0&gt;)
</code></pre>
<pre><code class="lang-python">def hook(self, grad_input, grad_output):
    print(grad_input)
    print(grad_output)

model = LinearTransformation(constant=True)

x = torch.tensor([1., 2.])

model.linear.register_backward_hook(hook)
model(x).backward()
# (tensor([60., 60.]), tensor([60., 60.]))
# (tensor([60., 60.]),)
</code></pre>
<p>My questions are</p>
<ul>
<li>For <code>register_forward_pre_hook</code> (first snippet), why 5, which is the final output, is also returned when I just register hook for <code>nn.Linear</code>.</li>
<li>For <code>register_backward_hook</code> (second snippet), I am not sure what these <code>tensor([60, 60])</code> correspond to. I could see <strong>maybe</strong> <code>grad_output</code> is gradient respect to output of <code>nn.Linear</code>. But how about another two <code>tensor([60, 60])</code>?</li>
</ul>

username: albanD
cooked: <p>Hi,</p>
<p>In the first case, only the input should be given, not the output.<br>
Are you running this in an interpreter? That would explain why the result of <code>model(x)</code> is printed in both your examples.</p>
<p>For the backward hook, as you can see in the documentation, they are not working as expected at the moment. So you should not use them <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

username: tzm1003306213
cooked: <p>Hi,</p>
<p>Can I use hook to add a parameter masking function to Conv2d. Specifically, I’d like to add binary mask buffer to each conv2d module, during each training step, I need to update the mask buffer and then use it to mask the weight.</p>
<p>Thanks!</p>

username: MrRobot
cooked: <p>I guess you could write the masking part in the model <code>forward()</code> function. Unless you would like finer control of the type of mask you use.</p>

username: albanD
cooked: <p>If you are using nightly build, we just landed a pruning utility in <code>nn.utils</code>. The upcoming tutorial for it can be found <a href="https://deploy-preview-605--pytorch-tutorials-preview.netlify.com/intermediate/pruning_tutorial.html?highlight=pruning">here</a>.</p>
<p>Otherwise, you will have to emulate what is done in the pruning module yourself.</p>

username: tzm1003306213
cooked: <p>Thank you so much! I’ll try it in the next days. Is this utility support multi-GPU training and mixed precision training?</p>

username: Michela
cooked: <p>Yes, for the use cases that have been tested, pruning works well with <code>DataParallel,</code> <code>DistributedDataParallel</code>, as well as <code>apex</code> and <code>torch.quantization</code>.<br>
Please flag anything that doesn’t work as you try out the functionality.</p>

username: mohit117
cooked: <p>So then what is the best way to check gradients for each layer? I used to apply a forward hook using ‘register_forward_hook’ on each layer and was thinking of doing the same checking gradients on each layer by using ‘register_backward_hook’ for each layer.</p>

username: albanD
cooked: <p>Hi,</p>
<p>The best way to do it is to use <code>register_hook()</code> on the Tensors you want the gradients for during the forward.</p>