How to modify Conv2d input gradients using backward hook?

braindotai: <p>I want to change the gradients during a backward pass for each Conv2d modules, so I’m trying to figure out how to change the input_gradiens using the backward hook, but cannot figure out what to return from the hook function in order to change the input_gradients.</p>
<pre><code class="lang-python">
def backward_hook(module, input_grads, output_grads):
        if isinstance(module, nn.Conv2d):
            input_grads = ..... # changing the input_grads here
            return input_grads, output_grads

    for module in model.modules():
        module.register_backward_hook(backward_hook)
</code></pre>
<p>And I’m getting this error on calling the <code>.backward()</code></p>
<pre><code class="lang-auto">TypeError: expected Variable, but hook returned 'tuple'
</code></pre>
<p>I’ve tried all permutations and combinations 🥲 such as:</p>
<pre><code class="lang-python">return input_grads[0], output_grads[0]
</code></pre>
<pre><code class="lang-python">return input_grads, output_grads[0]
</code></pre>
<pre><code class="lang-python">return input_grads,
</code></pre>
<pre><code class="lang-python">return input_grads[0]
</code></pre>
<p>…</p>
<p>etc but still no luck.</p>
<p>What should I return in the case of <code>Conv2d</code>? Also, I could not find any documentation related to this, so very difficult to get some results.<br>
I’ve also noticed that what we need to return in the backward hook function also depends on the type of the module, for instance the code below runs just fine with the <code>nn.ReLU</code>:</p>
<pre><code class="lang-python">
def backward_hook(module, input_grads, output_grads):
        if isinstance(module, nn.ReLU):
            input_grads = ..... # changing the input_grads here
            return input_grads, # so with relu it seems we don't need to return output_grads, 
            # but this does not work with Conv2d

    for module in model.modules():
        module.register_backward_hook(backward_hook)
</code></pre>

ptrblck: <p>I don’t think you should return the <code>output_grads</code>, as they were already passed to the module and won’t be used anymore. Instead you should return all input gradients, which would be passed to the “previous” layer (previous in the sense of the forward execution). Also use <code>register_full_backward_hook</code>, as <code>register_backward_hook</code> might not be working correctly.<br>
From the docs:</p>
<blockquote>
<p>The <code>grad_input</code> and <code>grad_output</code> are tuples that contain the gradients with respect to the inputs and outputs respectively. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the input that will be used in place of <code>grad_input</code> in subsequent computations. <code>grad_input</code> will only correspond to the inputs given as positional arguments and all kwarg arguments are ignored. Entries in <code>grad_input</code> and <code>grad_output</code> will be <code>None</code> for all non-Tensor arguments.</p>
</blockquote>
<p>Here is a minimal example:</p>
<pre data-code-wrap="python"><code class="lang-python">def backward_hook(m, input_gradients, output_gradients):
    print('input_gradients {}'.format(input_gradients))
    print('output_gradients {}'.format(output_gradients))
    input_gradients = (torch.ones_like(input_gradients[0]), )
    return input_gradients

conv = nn.Conv2d(1, 1, 3)
conv.register_full_backward_hook(backward_hook)

x = torch.randn(1, 1, 3, 3).requires_grad_()
out = conv(x)
out.mean().backward()
print(x.grad) # ones
</code></pre>

braindotai: <p>Thanks a lot <a class="mention" href="/u/ptrblck">@ptrblck</a>, always to the rescue <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=10" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

chaudatascience: <p>Hi,</p>
<p>I have a question regarding the shape of <code>input_grads</code> and <code>output_gradients</code>.<br>
I’ve checked on some models, but in all cases, <code>len(input_gradients)</code> and <code>len(output_gradients)</code>   are always 1. Is it the case for all input_gradients/output_gradients?<br>
If not, can you show me one example where input_gradients/output_gradients have more than 1 entry?<br>
Thanks!</p>

ptrblck: <p>No, that’s not always the case and depend on the number of input and output arguments.<br>
Here is a small example showing the usage of 3 input and 2 output arguments where one input argument is a constant:</p>
<pre><code class="lang-python">def backward_hook(m, input_gradients, output_gradients):
    print('input_gradients {}'.format(input_gradients))
    print('output_gradients {}'.format(output_gradients))
    return input_gradients

class MyModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(1, 1)
        self.fc2 = nn.Linear(2, 2)
        
    def forward(self, x, y, static):
        x = self.fc1(x) + static
        y = self.fc2(y)
        return x, y

model = MyModule()
model.register_full_backward_hook(backward_hook)

x = torch.randn(1, 1).requires_grad_()
y = torch.randn(1, 2).requires_grad_()
static = 1.
out_x, out_y = model(x, y, static)
loss = out_x + out_y
loss.mean().backward()
# input_gradients (tensor([[-0.2857]]), tensor([[-0.1686,  0.3697]]), None)
# output_gradients (tensor([[1.]]), tensor([[0.5000, 0.5000]]))
</code></pre>
<p>and thus its <code>input_gradient</code> is <code>None</code>.<br>
You can play around with this code snippet and add more input/output arguments.</p>

chaudatascience: <p>This makes sense. Thanks a lot!</p>