Understanding gradient calculation with backward_pre_hooks

Ricu: <p>Hello everyone.</p>
<p>A little time ago I already <a href="https://discuss.pytorch.org/t/linear-layer-can-not-register-backward-pre-hook/175379">asked</a> about how to use the <a href="https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.register_full_backward_pre_hook" rel="noopener nofollow ugc">full_backward_pre_hooks</a>.</p>
<p>From my understanding, these can be used to manipulate the calculation of the backward pass of the affected layer.</p>
<p>We have three components:</p>
<ul>
<li>grad_output (coming from the output of the model)(accesible via the <a href="https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.register_full_backward_hook" rel="noopener nofollow ugc">full_backward_hook</a> and <a href="https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.register_full_backward_pre_hook" rel="noopener nofollow ugc">full_backward_pre_hook</a>)</li>
<li>grad_input (going towards the input of the model)(accesible via the <a href="https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.register_full_backward_hook" rel="noopener nofollow ugc">full_backward_hook</a> , but NOT the <a href="https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.register_full_backward_pre_hook" rel="noopener nofollow ugc">full_backward_pre_hook</a>)</li>
<li>grad (of the current layer, so the layers gradient of the form dLayer_output/dLayer_input)(stored in the parameters <a href="https://pytorch.org/docs/stable/generated/torch.Tensor.grad.html" rel="noopener nofollow ugc">.grad()</a>)</li>
</ul>
<p>So the calculation is approximately as follows:</p>
<pre><code class="lang-auto">grad_input = grad_output*grad
</code></pre>
<p>So to replace the incoming grad_ouput of a layer and therefore modify the following backward pass computations (in particular the grad_input), I expected using the <a href="https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.register_full_backward_pre_hook" rel="noopener nofollow ugc">full_backward_pre_hook</a> would would be the solution. But runnning the following code, doesn’t confirm my intuition:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn

class Backward_Debug_Hook():
  def __init__(self, module):
    self.hook = module.register_full_backward_hook(self.hook_fn)

  def hook_fn(self, module, grad_input, grad_output):
    print('grad_output')
    print(grad_output)
    print('grad_input')
    print(grad_input)
    
  def close(self):
    self.hook.remove() 

class Insert_Hook():
  def __init__(self, module, new_grad_output=None):
    self.new_grad_output = new_grad_output
    self.hook = module.register_full_backward_pre_hook(self.hook_fn)

  def hook_fn(self, module, grad_output):
    return self.new_grad_output

  def close(self):
    self.hook.remove()

# simple model
model = nn.Sequential(
  nn.Linear(2, 2),
  nn.Sigmoid(),
  nn.Linear(2,2)
)
last_layer = model[-1]

debug_hook = Backward_Debug_Hook(last_layer) # attach debug hook
x = torch.randn(1, 2) # artificial input
out = model(x) # forward pass
print('without gradient insertion')
out.mean().backward() # backward pass


model.zero_grad()

artifical_grad = (100*torch.ones([1,2]),)
insert_hook = Insert_Hook(last_layer,artifical_grad)
out = model(x) # forward pass
print('with gradient insertion')
out.mean().backward() # backward pass
</code></pre>
<p>as the output is.</p>
<pre><code class="lang-auto">grad_output
(tensor([[0.5000, 0.5000]]),)
grad_input
(tensor([[ 0.3774, -0.1604]]),)
with gradient insertion
grad_output
(tensor([[100., 100.]]),)
grad_input
(tensor([[ 0.3774, -0.1604]]),)
</code></pre>
<p>According to my understanding from above, I expected the grad_input to be different the second time.</p>
<p>Where lies the mistake in my understanding?</p>

Ricu: <p>Bump for visibility.</p>

soulitzer: <p>Ah this is probably a bug.</p>
<pre><code class="lang-python">import torch
import torch.nn as nn

a = torch.ones(2, requires_grad=True)

model = nn.Linear(2, 2)

def fn(module, grad_output):
    return (grad_output[0] * 0,)

model.register_full_backward_pre_hook(fn, prepend=False)

out = model(a)
out.sum().backward()
print(a.grad) # should be 0, but its not
</code></pre>

soulitzer: <p>It should be fixed in a nightly in the next couple days, but for now the following fix should be applied:</p>
<pre><code class="lang-auto">diff --git a/torch/utils/hooks.py b/torch/utils/hooks.py
index 6d5a97d4288e4..2cc7efeb5c124 100644
--- a/torch/utils/hooks.py
+++ b/torch/utils/hooks.py
@@ -223,6 +223,10 @@ def hook(_, grad_output):
                             raise RuntimeError("Backward hook for Modules where no input requires "
                                                "gradient should always return None or None for all gradients.")
                     self.grad_outputs = None
+
+                if self.grad_outputs is not None:
+                    return tuple(self.grad_outputs[i] for i in self.output_tensors_index)
+
             grad_fn.register_hook(hook)
 
         is_tuple = True
</code></pre>

Ricu: <p>Thanks for your reply. Glad my understanding wasn’t wrong after all.</p>
<p>Hope this gets fixed quickly in a stable release soon.</p>
<p>Thanks also for the quick fix. Will try later and notfiy if its still not working. Cheers</p>