How can I modify a ReLU layer's backward?

tumble-weed: <p>I want to modify the backward of relu, such that i simply pass through the gradients coming from the top rather than 0-ing out the ones where the unit is off.</p>
<p>I looked at <a href="https://discuss.pytorch.org/t/exact-meaning-of-grad-input-and-grad-output/14186/5">this thread</a> and couldn’t get much out of it. Looking <a href="https://discuss.pytorch.org/t/manipulating-gradients-in-backward/15487">here</a>,  while this works for making the gradients zero, i am not sure what the components of <code>gradient_input</code> are and which one i should modify to have a pass-through.</p>

albanD: <p>You can use the backward hooks on the relu Module:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn

inp = torch.randn(5, requires_grad=True)
print(inp) # [ 1.7485,  0.5180, -1.2110, -1.3865, -0.6293]

relu = nn.ReLU()

# Without hook:
out = relu(inp)
out.sum().backward()
print(inp.grad) # [1., 1., 0., 0., 0.]

inp.grad.zero_()
# With hook
def grad_hook(mod, grad_in, grad_out):
    # Ignore the computed grad_input
    # Make it look like this is an identity
    return grad_out
relu.register_backward_hook(grad_hook)

out = relu(inp)
out.sum().backward()
print(inp.grad) # [1., 1., 1., 1., 1.]
</code></pre>

tumble-weed: <p>thanks for the reply, from <a href="https://discuss.pytorch.org/t/manipulating-gradients-in-backward/15487/2?u=tumble-weed">this post</a> however it seems like what one has to return is modification of <code>grad_input</code>, where in the snippet in your post, we seem to be returning a modification of <code>grad_out</code>.</p>

albanD: <p>That depends what you want to do.<br>
If you want your relu’s backward to act as if it was an identity, then the gradient of the output wrt to input is 1, and so what you want your backward to compute is <code>grad_output * 1</code>. Ignoring the grad_input computed by the backward of the relu as you replace it.</p>

tumble-weed: <p>yeah i see your point, for a pointwise activation <code>grad_input</code> and <code>grad_output</code> would be of the same size, so its straightforward to directly use the <code>grad_output</code> to substitute <code>grad_input</code>.</p>
<p>An extra question please, form other posts it seems that modification of gradients is not so straightforward for other layers, is that still the case?</p>

fourleafclover: <p>Hi, thanks a lot for answering this question <a class="mention" href="/u/alband">@albanD</a>. Your example already makes it quite clear for me, however, I do have an additional question just to clarify it completely:<br>
While searching the web, I was not really able to find a clear documentation on the parameters <code>grad_input</code> and <code>grad_output</code> of the “hook function” that is passed to <code>register_backward_hook</code>, beyond what is stated <a href="https://pytorch.org/docs/stable/generated/torch.nn.modules.module.register_module_full_backward_hook.html#torch.nn.modules.module.register_module_full_backward_hook" rel="noopener nofollow ugc">here</a>. This is how I understand it based on your little code example above:</p>
<p>In the backward pass, something like this is computed:</p>
<p>Let x be the randomly initialized tensor and  u = ReLU(x)</p>
<p>Then we would have something like:</p>
<p>dL/dx = du/dx * dL/du (by the chain rule)</p>
<p>where L is the “loss” (i.e. the sum of the elements of u)</p>
<p>In this case du/dx should represent <code>grad_input</code> (i.e. the gradient of the layer output w.r.t. the layer input) and dL/du should represent <code>grad_output</code> (i.e. the gradient of the loss w.r.t. the layer output).</p>
<p>Is my interpretation correct?</p>
<p>(PS: Sorry for the math formatting, Latex is not supported. I hope it still becomes clear.)</p>