Re-register a removed backward hook within the backward function

jerrybai1995: <p>Hi!</p>
<p>I wonder if there is a way to re-register the same backward hook handle after we have removed it? Specifically, my layer’s custom backward function needs to use the Jacobian information of the output, and one way of doing this is:</p>
<pre><code class="lang-python">class MyLayer(nn.Module):
        def __init__(self, ...):
                ...  
                self.hook = None

	def forward(x):
		x.require_grad_()
		y = self.net(x)

		def backward_hook(grad):
			if self.hook is not None:
				self.hook.remove()
			return autograd.grad(y, x, grad)[0]

		self.hook = y.register_hook(backward_hook)
		return y
</code></pre>
<p>In the code above, the <code>self.hook.remove()</code> is needed because otherwise the <code>autograd.grad(y,x,...)</code> would recursively call the <code>backward_hook</code> function. When I have one loss per training iteration and backprop through it, everything works fine.</p>
<p>However, one obvious drawback of the implementation above is that if I have <strong>two losses</strong>, <code>loss1(y)</code> and <code>loss2(y)</code> that <em>both</em> depends on the output <code>y</code>, and if I do</p>
<pre><code class="lang-auto">loss1(y).backward()
loss2(y).backward()
</code></pre>
<p>the second loss’ (i.e., <code>loss2</code>) backward pass will not go through the <code>backward_hook</code> function because I have already removed the handle when I call <code>loss1(y).backward()</code>. And that is not an expected behavior. (For example, if I do <code>gradcheck</code>, this problem will occur.)</p>
<p>One way to work around this issue is by cloning and detaching the <code>x</code> and <code>y</code> in the layer; i.e.,</p>
<pre><code class="lang-python">class MyLayer(nn.Module):
	def forward(x):
		y = self.net(x)

		xc = x.clone().detach().requires_grad_()
		yc = self.net(xc)
		def backward_hook(grad):
			return autograd.grad(yc, xc, grad)[0]

		y.register_hook(backward_hook)
		return y
</code></pre>
<p>But this requires calling <code>self.net(xc)</code> again, which adds computational and memory costs. I wonder if there’s a way that we can re-register the hook <strong>after</strong> we have already removed it (as in the first code block), so that the custom backward pass works even if we call it multiple times, and we don’t need to clone anything?</p>
<p>Thanks!</p>

albanD: <p>Hi,</p>
<p>I don’t think I can thing of a “safe” way to do that. Mainly because you will have to change the state of the Module while the backward is running to avoid the infinite recursion.<br>
So calling two backward concurrently will always do something bad.</p>
<p>Assuming that you never do concurrent backwards, you can set a flag on the Module to know if this is the “top” backward or not:</p>
<pre><code class="lang-python">class MyLayer(nn.Module):
        def __init__(self, ...):
                ...  
                self.in_backward = False

	def forward(x):
		x.require_grad_()
		y = self.net(x)

		def backward_hook(grad):
			if self.in_backward:
                # If we're already in backward, do nothing
				return grad
			return autograd.grad(y, x, grad)[0]

		y.register_hook(backward_hook)
		return y
</code></pre>