How to avoid recalculating a function when we need to backpropagate through it twice?

username: sliorde
cooked: <p>I want to do the following calculation:</p>
<pre><code class="lang-python">l1 = f(x.detach(), y)
l1.backward()
l2 = -1*f(x, y.detach())
l2.backward()
</code></pre>
<p>where <code>f</code> is some function, and <code>x</code> and <code>y</code> are tensors that require gradient. Notice that <code>x</code> and <code>y</code> may both be the results of previous calculations which utilize shared parameters (for example, maybe <code>x=g(z)</code> and <code>y=g(w)</code> where <code>g</code> is an <code>nn.Module</code> ).</p>
<p>The issue is that <code>l1</code> and <code>l2</code> are both numerically identical, up to the minus sign, and it seems wasteful to repeat the calculation <code>f(x,y)</code> twice. It would be nicer to be able to calculate it once, and apply <code>backward</code> twice on the result. Is there any way of doing this?</p>
<p>One possibility is to manually call <code>autograd.grad</code> and update the <code>w.grad</code> field of each <code>nn.Parameter</code> <code>w</code> . But I’m wondering if there is a more direct and clean way to do this, using the <code>backward</code> function.</p>

username: albanD
cooked: <p>Hi,</p>
<p>I’m afraid the only thing you can do is to give <code>retain_graph=True</code> the first time you call backward.</p>

username: sliorde
cooked: <p>Can you please elaborate on how this would help?</p>

username: albanD
cooked: <p>This flag prevents the graph internal states from being freed. So you will be able to backprop through the graph as many times as you want (as mentioned in the error message you shared).</p>

username: sliorde
cooked: <p>Thanks. I would appreciate if you can give a code snippet demonstrating your suggestion. I am aware of  <code>retain_graph</code>, but I don’t see how setting it to <code>True</code> will help accomplish what I asked for. It seems to me that there is still a missing ingredient that would let me to do <code>backward</code> twice on the retained graph, each time keeping a different part of the graph detached.</p>

username: albanD
cooked: <pre><code class="lang-auto">l1 = f(x.detach(), y)
l1.backward(retain_graph=True)
l2 = -1*f(x, y.detach())
l2.backward()
</code></pre>
<p>This should not throw the error you mentioned <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

username: sliorde
cooked: <p>Thank you, but I did not mention any error. I wanted to avoid calculating the function <code>f</code> twice. I think that in your solution, it is calculated twice, right? Unless PyTorch automatically caches the first calculation of <code>f</code> and then in the second invocation the cached graph is reused somehow, including the detached subgraph composed of the ancestors of <code>x</code>?</p>

username: albanD
cooked: <p>Ho sorry I think I misread the title of the topic and though it was an error <img src="https://discuss.pytorch.org/images/emoji/apple/frowning.png?v=9" title=":frowning:" class="emoji" alt=":frowning:"></p>
<p>Does the function <code>f</code> has any parameter into it?</p>
<p>If there is nothing else in there, and the <strong>only</strong> way to get to the parameters is via x and y, I would do:</p>
<pre><code class="lang-python">x, y = g(input, params)
# f must have NO parameters
# Equivalent (by linearity of the gradient) to 
# l1 = f(x.detach(), y)
# l1.backward()
# l2 = -1*f(x, y.detach())
# l2.backward()
y.register_hook(lambda x: -x)
loss = f(x, y)
loss.backward()
</code></pre>
<p>Note that you should double check that you get the same gradients still as I might be missing something <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

username: sliorde
cooked: <p>This solved my problem nicely. Thanks.<br>
Here is code demonstrating that the solution works:</p>
<pre><code class="lang-python">import torch
lin = torch.nn.Linear(1,1,bias=False)
lin.weight.data[:] = 1.0
a = torch.tensor([1.0])
b = torch.tensor([2.0])
loss_func = lambda x,y: (x-y).abs()

# option 1: this is the inefficient option, presented in the original question
lin.zero_grad()
x = lin(a)
y = lin(b)
loss1 = loss_func(x.detach(),y)
loss1.backward(retain_graph=True)
loss2 = -1*loss_func(x,y.detach()) # second invocation of `loss_func` - not efficient!
loss2.backward()
print(lin.weight.grad)

# option 2: this is the efficient method, suggested by @albanD. 
lin.zero_grad()
x = lin(a)
y = lin(b)
x.register_hook(lambda t: -t)
loss = loss_func(x,y) # only one invocation of `loss_func` - more efficient!
loss.backward()
print(lin.weight.grad) # the output of this is identical to the previous print, which confirms the method

# option 3 - this should not be equivalent to the previous options, used just for comparison
lin.zero_grad()
x = lin(a)
y = lin(b)
loss = loss_func(x,y)
loss.backward()
print(lin.weight.grad)

</code></pre>

username: sliorde
cooked: <p>However, it would be nice to know if there is a more general alternative, that works also for non-linear cases. In other words, it would be interesting to know how to perform the following calculation, but reducing to a single invocation of <code>f(x,y)</code>, instead of two:</p>
<pre><code class="lang-python">l1 = u(f(x.detach(), y))
l1.backward()
l2 = v(f(x, y.detach()))
l2.backward()
</code></pre>
<p>Here, <code>u</code> and <code>v</code> are two functions.</p>

username: albanD
cooked: <p>Hi,</p>
<blockquote>
<p>non-linear cases</p>
</blockquote>
<p>The gradient is always linear <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"><br>
This won’t work if u/v/f have paramters though because the flipping of the gradient for the second loss only happens before <code>f</code>. So parameters in these functions would just see the sum of the two losses.</p>