Linear layer can not register backward pre hook

Ricu: <p>Hello, I am using the following versions:</p>
<p>PyTorch Version:  1.13.1+cu116<br>
Torchvision Version:  0.14.1+cu116</p>
<p>I am trying to insert a backward pre hook into a nn.Linear layer:</p>
<pre><code class="lang-auto">class Insert_Hook():
  def __init__(self, module, new_grad_output):
    self.new_grad_output = new_grad_output
    # use prepend=True so that this is definetly the first hook being applied
    self.hook = module.register_full_backward_pre_hook(self.hook_fn)

  def hook_fn(self, module, grad_input, grad_output):
    # simply return the previously caught grad_output
    # this will replace the current grad_output
    return self.new_grad_output
  
  def close(self):
    self.hook.remove()

artifical_grad = (torch.ones([3,2]),)
print(artifical_grad)
insert_hook = Insert_Hook(affected_layer,artifical_grad)
</code></pre>
<p>Here, affected_layer is a nn.Linear layer of the correct size.</p>
<p>It gives me the the error</p>
<pre><code class="lang-auto">'Linear' object has no attribute 'register_full_backward_pre_hook'
</code></pre>
<p>This surprises me  as I thought nn.Linear subclasses nn.Module. The latter one should definetly have the method as described here <a href="https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.register_full_backward_pre_hook" rel="noopener nofollow ugc">Reference</a>.</p>
<p>EDIT:<br>
It will also not let me use the prepend statement:</p>
<pre><code class="lang-auto">register_full_backward_hook() got an unexpected keyword argument 'prepend'
</code></pre>

ptrblck: <p>You are linking to the docs of the current <code>2.0.0</code> release, while you are using <code>1.13.1</code> which doesn’t seem to provide this method as seen <a>here</a>.<br>
Update PyTorch and you should be able to use <code>register_full_backward_pre_hook</code>.</p>

Ricu: <p>Okay, rookie mistake on my side. Still, I updated my pytorch to v2.0.0 but I still get the same message.</p>
<p>I noticed that the <a href="https://pytorch.org/get-started/pytorch-2.0/" rel="noopener nofollow ugc">update site</a>, says:</p>
<blockquote>
<p>Module and Tensor <a href="https://pytorch.org/docs/stable/notes/modules.html#module-hooks" rel="noopener nofollow ugc">hooks</a> don’t fully work at the moment, but they will eventually work as we finish development.</p>
</blockquote>
<p>So is this the problem and backward_pre_hooks are not working right now?</p>

ptrblck: <p>I think the warning is related to the newly introduced <code>torch.compile</code> backend not eager mode as it’s working for me in <code>2.0.0</code>:</p>
<pre><code class="lang-python">model = nn.Linear(10, 10)

def get_hook(name):
    def hook(module, grad_output):
        print("module {} called with grad_output.abs().sum() {}".format(
            name, grad_output[0].abs().sum()))
    return hook

model.register_full_backward_pre_hook(get_hook("linear"))

x = torch.randn(1, 10)
out = model(x)
out.mean().backward()
# module linear called with grad_output.abs().sum() 1.0000001192092896
</code></pre>

Ricu: <p>Sorry, I am a little confused by your wording.</p>
<p>How is this connected to the new torch.compile() function?</p>

ptrblck: <p>The <code>Hooks</code> section claiming:</p>
<blockquote>
<p>Module and Tensor <a href="https://pytorch.org/docs/stable/notes/modules.html#module-hooks">hooks</a> don’t fully work at the moment, but they will eventually work as we finish development.</p>
</blockquote>
<p>is only referring to the <code>torch.compile</code> mode, not any other usage as seen in my code snippet.</p>

Ricu: <p>I finally found the mistake and it was very stupid. I used</p>
<pre><code class="lang-auto">layer.register_full_backward__pre_hook(..)
</code></pre>
<p>instead of</p>
<pre><code class="lang-auto">layer.register_full_backward_pre_hook(..)
</code></pre>
<p>(so just an extra underscore).</p>
<p>Still, thank you for your advice. It helped me as I wasn’t aware of the 2.0 update before.</p>
<p>Cheers</p>

Zeyuan.Yin: <p>I want to confirm that the function registered with <code>register_full_backward_pre_hook</code> is called before the backward computation of the module. And we can modify the <code>grad_output </code>before calculating the <code>grad_in</code>. Right?</p>
<p>And why I can’t find the <code>register_full_backward_pre_hook</code> in pytorch doc v2.0.0</p>

soulitzer: <aside class="quote no-group" data-username="Zeyuan.Yin" data-post="8" data-topic="175379">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/zeyuan.yin/48/54711_2.png" class="avatar"> Zeyuan.Yin:</div>
<blockquote>
<p>register_full_backward_pre_hook</p>
</blockquote>
</aside>
<p>You should use <a href="https://pytorch.org/docs/stable/generated/torch.nn.modules.module.register_module_full_backward_hook.html" class="inline-onebox" rel="noopener nofollow ugc">torch.nn.modules.module.register_module_full_backward_hook — PyTorch 2.1 documentation</a>. That should do what you describe.</p>

Zeyuan.Yin: <p>Thanks.<br>
But this is the description of <code>register_module_full_backward_hook</code>, which is used after backward<br>
I look for the description of <code>register_full_backward_pre_hook</code>, which is used before backward</p>

Zeyuan.Yin: <p>I give an example to modify the grad_output successfully.</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn

model = nn.Linear(10, 10)

def get_hook(name):
    def hook(module, grad_output):
        print("Before modify: module {} called with grad_output.abs().sum() {}".format(
            name, grad_output[0].abs().sum()))
        new_grad_output = grad_output[0].clone()*10

        return tuple([new_grad_output])

    return hook


def check_hook(name):
    def hook(module, grad_input, grad_output):
        print("After modify: module {} called with grad_output.abs().sum() {}".format(
            name, grad_output[0].abs().sum()))
    return hook

model.register_full_backward_pre_hook(get_hook("linear"))
model.register_full_backward_hook(check_hook("linear"))

x = torch.randn(1, 10)
out = model(x)
out.mean().backward()
# Output:
# Before modify: module linear called with grad_output.abs().sum() 1.0000001192092896
# After modify: module linear called with grad_output.abs().sum() 10.0
</code></pre>

soulitzer: <aside class="quote no-group" data-username="Zeyuan.Yin" data-post="8" data-topic="175379">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/zeyuan.yin/48/54711_2.png" class="avatar"> Zeyuan.Yin:</div>
<blockquote>
<p>register_full_backward_pre_hook</p>
</blockquote>
</aside>
<p>Okay I found it <a href="https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.register_full_backward_pre_hook" class="inline-onebox" rel="noopener nofollow ugc">Module — PyTorch 2.1 documentation</a>, the docs do exist but it wasn’t appearing in search results for some reason.</p>

soulitzer: <p><em>Edit: for anyone reading in the future, module backward pre hook is introduced in the 2.0.0 release but is buggy in that version, so if you want to use module backward pre hooks you should be using 2.0.1</em></p>
<p>FYI there’s a bug in backward pre hook. You should either update to a nightly in a couple days, or apply the patch in <a href="https://discuss.pytorch.org/t/understanding-gradient-calculation-with-backward-pre-hooks/176154/3" class="inline-onebox">Understanding gradient calculation with backward_pre_hooks - #3 by soulitzer</a> if you’d like to use module backward pre hooks.</p>
<p>Another workaround is to apply a hook to the output of your module instead.</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(grad_output):
    return grad_output * 2

def fn2(module, grad_inputs, grad_output):
    # The modification is still observed
    print(grad_inputs, grad_output)
    return (grad_inputs[0] / 2,)

# No longer using full backward pre hooks
# model.register_full_backward_pre_hook(fn)
model.register_full_backward_hook(fn2)

out = model(a)
# Instead, register a hook to the output of your module
# Unlike module hooks, register AFTER the forward runs
out.register_hook(fn)

out.sum().backward()
print(a.grad)
</code></pre>