Is it possible to use `register_backward_full_hook` on a different loss than the one used to compute the gradient?

username: AlphaBetaGamma96
cooked: <p>Hi All,</p>
<p>I have a quick question regarding <code>register_backward_full_hook</code>. To give a quick explanation on my model I effectively have a model which is a Feed-Forward Network that has <code>N</code> inputs and <code>L</code> number of <code>nn.Linear</code> layers. The output of my network is the sign of the output along with its log-absolute value, which is done via <code>torch.linalg.slogdet</code>.</p>
<p>Now, my loss function is effectively broken down into 2 parts. The first part calculates a scaling factor (can be positive or negative) which is detached so it holds no gradient, it only scales each sample of input data. Let’s define that as <code>scale_factor</code>, which will have the dimensions of <code>[B]</code> where <code>B</code> is the number of samples within my batch. The second part is the log-absolute value of the network, so the output from the <code>torch.linalg.slogdet</code> function but ignoring the sign part, so, <code>torch.linalg.slogdet(x)[1]</code>. I then define my loss to be the element wise product of these two losses, then mean reduced and that’s my loss, an example of this would be,</p>
<pre><code class="lang-auto">input_data = torch.randn(B,N) #input data is shape [B,N]

scale_factor = loss1(input_data) #returns Tensor of shape [B]

#net returns torch.linalg.slogdet, so just grab logabs value 
logabs_net = net(input_data)[1] 

loss = torch.mean( scale_factor.detach() * logabs_net )

optim.zero_grad()
loss.backward() #calculate gradients and call backward hooks (all in one go)
optim.step()
</code></pre>
<p>Within my network, I register <code>forward_pre_hook</code> and <code>backward_full_hook</code> on my Linear layers (because I seek to precondition my gradients via this information). However, the <code>grad_output</code> property I need is slightly different than the one PyTorch returns to me and I was wondering if you can place <code>backward_full_hooks</code> on a different loss value than the same one that’s used to calculate the gradients of your given loss.</p>
<p>As it currently stands the <code>grad_output</code> Tensor that’s return via the <code>backward_full_hook</code> returns the gradient of <code>scale_factor.detach() * logabs_net</code> with respect to the output of a given <code>nn.Linear</code> layer for <strong>all</strong> input samples.</p>
<p><strong>My question</strong>: Is it at all possible to change the hook such that it returns the gradient of <code>logabs_net</code> with respect to the output of a given <code>nn.Linear</code> layer? I’ve tried taking the returned value of the hook and dividing by the <code>scale_factor.detach()</code> Tensor. However, this value can be equal to 0, and in the case it is it crashes my code because I’m dividing by 0 when in fact the value of  the gradient of <code>logabs_net</code> with respect to output of a Layer will no non-zero and finite.</p>
<p>Any help on this would be greatly appreciated! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=10" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

username: AlphaBetaGamma96
cooked: <p>You can do this via using both loss function and just caching the needed values</p>