Status of register_backward_hook

username: jamesproud
cooked: <p>It appears that backward hooks are currently broken as the recent issue and PR below discuss.</p>
<p><aside class="onebox githubissue">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/issues/12331" target="_blank" rel="nofollow noopener">github.com/pytorch/pytorch</a>
  </header>
  <article class="onebox-body">
    <a href="https://github.com/ezyang" rel="nofollow noopener">
<img src="https://avatars1.githubusercontent.com/u/13564?v=2&amp;s=96" class="thumbnail onebox-avatar" width="96" height="96">
</a>

<h4><a href="https://github.com/pytorch/pytorch/issues/12331" target="_blank" rel="nofollow noopener">Issue: Feedback about PyTorch register_backward_hook</a></h4>

<div class="date" style="margin-top:10px;">
	<div class="user" style="margin-top:10px;">
	opened by <a href="https://github.com/ezyang" target="_blank" rel="nofollow noopener">ezyang</a>
	on <a href="https://github.com/pytorch/pytorch/issues/12331" target="_blank" rel="nofollow noopener">2018-10-04</a>
	</div>
	<div class="user">
	</div>
</div>

<pre class="content" style="white-space: pre-wrap;">Maxim Naumov @mnaumovfb sends in the following feedback:
I thought I would share some feedback from my recent work with register_backward_hook function....</pre>

<div class="labels">
</div>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>
<br>
<aside class="onebox githubpullrequest">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/pull/12573" target="_blank" rel="nofollow noopener">github.com/pytorch/pytorch</a>
  </header>
  <article class="onebox-body">
      <a href="https://github.com/albanD" target="_blank" rel="nofollow noopener">
    <img alt="albanD" src="https://avatars3.githubusercontent.com/u/6359743?v=4" class="thumbnail onebox-avatar" width="90" height="90">
  </a>

<h4>
  <a href="https://github.com/pytorch/pytorch/pull/12573" target="_blank" rel="nofollow noopener">nn.Module hook fix and improvements</a>
</h4>

<div class="date">
  by <a href="https://github.com/albanD" target="_blank" rel="nofollow noopener">albanD</a>
  on <a href="https://github.com/pytorch/pytorch/pull/12573" target="_blank" rel="nofollow noopener">06:09PM - 11 Oct 18 UTC</a>
</div>

<div class="github-commit-stats">
  <strong>6 commits</strong>
  changed <strong>5 files</strong>
  with <strong>193 additions</strong>
  and <strong>61 deletions</strong>.
</div>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>
</p>
<p>Is there a path towards fixing or appropriate operation for implementing a backward operation in lieu of using <code>register_backward_hook</code>?</p>

username: albanD
cooked: <p>Hi,</p>
<p>Unfortunately it is a really tricky problem to solve without impacting users not using the hook and supporting most of the <code>nn.Module</code> behaviours.<br>
It is very hard to make it work nicely with inplace operations and not even possible at the moment to forbid inplace ops if we don’t support it.<br>
So there is some progress in cleaning this inplace ops and the backward hooks will come after.</p>
<p>Unfortunately, I have limited amount of time to do all the changes that are needed but if someone wants to work on this I can help.</p>
<p>At the moment, the workaround is to use <code>register_hook</code> directly on the specific Tensors that you want to change.</p>

username: jamesproud
cooked: <p>For my use case, using <code>register_hook</code> would work as there is only one layer of gradients that I want to target and modify on the backward pass. What’s the best way of implementing this?</p>

username: albanD
cooked: <p>Hi,</p>
<p>Could you give more details (or a small code sample) of what you want to do exactly? That way I can give you an exact code sample.</p>

username: jamesproud
cooked: <p>I’ve included some sample code below. I am trying to mask the gradients of <code>layer1</code> in this example. As you can see, I’m zeroing out the weights upon initialisation.</p>
<pre><code class="lang-auto">class MaskedLinear(nn.Module):
    # Currently unused. Intended for backward hook.
    def _zero_grad_mask(self, module, grad_input, grad_output):
        new_grad = Variable(
            torch.Tensor(grad_input[0].shape).cuda().masked_scatter_(self.indices_mask_tensor, grad_input[0].data)
        )
        return (grad_input[0], None, new_grad)

    def __init__(self, in_features, out_features, indices_mask):
        super(MaskedLinear, self).__init__()
        self.linear = nn.Linear(in_features, out_features).cuda()
        good_weights = self.linear.weight.data[indices_mask]
        self.linear.weight.data[:] = 0 # zero all weights out
        self.linear.weight.data[indices_mask] = good_weights # populate good weights
        self.indices_mask_tensor = torch.ByteTensor(out_features, in_features).cuda()
        self.indices_mask_tensor[indices_mask] = 1

    def forward(self, input):
        # Repeating original masking here is currently used, and works as a hack, but is incredibly slow.
        good_weights = self.linear.weight.data[self.indices_mask]
        self.linear.weight.data[:] = 0
        self.linear.weight.data[self.indices_mask] = good_weights
        return self.linear(input)

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.layer1 = MaskedLinear(in_dim, out_dim, mask)
        
    def forward(self, x):
        x = F.relu(self.layer1(x))
        return x
</code></pre>

username: albanD
cooked: <p>Hi,</p>
<p>In that case, you can register the hook directly on <code>self.linear.weight</code>. See the sample below:</p>
<pre><code class="lang-auto">import torch
from torch import nn

l = nn.Linear(5, 5)

inp = torch.rand(3, 5)

print("original weights")
print(l.weight)

opt = torch.optim.SGD(l.parameters(), lr=0.01)

mask = torch.rand(5, 5).gt(0.5)
print("our mask")
print(mask)
def hook_fn(grad):
    # You are not allowed to modify inplace what is given !
    out = grad.clone()
    out[mask] = 0
    return out

# Registering the hook only once. This can be done in
# MaskedLinear's __init__ function
l.weight.register_hook(hook_fn)

print("Loss goes down")
for _ in range(50):
    loss = l(inp).abs().sum()
    opt.zero_grad()
    loss.backward()
    opt.step()
    print(loss)
print("only some weights have changed")
print(l.weight)
</code></pre>

username: Xiao_Song
cooked: <p>Hi,</p>
<p>I have a complex module where I want to have gradient back prop over a specific region of input. Can I do this in the forward function on the input?</p>
<pre><code class="lang-auto">def forward(self, input):
    self.loss = nn.MSELoss(input, self.target)
   
    def hook_fn(grad):
         return grad * self.mask // where mask is a binary mask defined to be a subset of input 
    input.register_hook(hook_fn)
   
    return input
</code></pre>

username: albanD
cooked: <p>Hi,</p>
<p>Yes registering a hook on the Tensor of interest during the forward is the way to go.</p>

username: Gaoyuan_Yang
cooked: <p>Hi, thanks for your reply. I have a question inside the hook_fn.<br>
Why you use grad.clone() inside the hook function, can’t we just use</p>
<pre><code class="lang-python">def hook_fn(grad):
    grad[mask] = 0
    return grad
</code></pre>
<p>I notice your comment “you are not allowed to modify inplace what is given”<br>
Does this mean we need to clone gradient inside register_hook functions for tensors?<br>
Since the tow functions have the same return values.<br>
Thanks a lot.</p>

username: albanD
cooked: <p>Hi,</p>
<p>Yes you have to clone it if you plan on doing inplace changes!<br>
This grad Tensor might be used in other places in the autograd and changing its value inplace can lead to silently wrong gradients!</p>