Trying to backward through the graph a second time in a loop

Jazzy66: <p>Hi, guys, I am new of using torch and such error gets raised in a simple loop.</p>
<p>To be concise, the loop aims to compute the gradient of a loss function, L, with respect to a vector, V, and update V till convergence. Here is a simple fake code helps to understand my problem.</p>
<pre><code class="lang-auto">import torch
import torch.nn.functional as F
import numpy as np

# functions
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

def argmax(x, axis=-1):
    return F.one_hot(torch.argmax(x, dim=axis), list(x.shape)[axis]).float()

def gradient(I,U,X):
    del X.grad
    Qp        =  argmax(U + X)
    gVp       =  torch.inner(Qp, U) + torch.inner(Qp, X)
    gVp_norm  =  torch.linalg.norm(gVp)
    gVp_norm.backward(torch.ones_like(gVp_norm))

    return (torch.inner(I,X) - gVp) * (I - X.grad)

# parameters
beta         =  0.98 
alpha        =  0.03
delta        =  0.1
T            =  10
kss          =  ((1 / beta - (1 - delta)) / alpha)**(1 / (alpha - 1))
k            =  np.linspace(0.5 * kss, 1.8 * kss, T)
k_reshaped   =  k.reshape(-1, 1)
c            =  k_reshaped ** alpha + (1 - delta) * k_reshaped - k

# iterations
V            =  torch.zeros(T, dtype = torch.float32, device = device, requires_grad = True)
VS           =  V
U            =  torch.tensor(c, dtype = torch.float32, device = device, requires_grad = True)
I            =  torch.eye(T, dtype = torch.float32, device = device, requires_grad = True)

for i in range(T):
    V_grad   =  gradient(I[i,:],U[i,:],V)
    VS_grad  =  gradient(I[i,:],U[i,:],VS)

    V  =  V - 0.1 * (V_grad - VS_grad)
</code></pre>
<p>I find that the error is most likely caused by the last line of the code <code>V  =  V - 0.1 * (V_grad - VS_grad)</code>, maybe it’s because the vector V being reused to compute the gradient multiple times. I tried <code>retain_graph=True</code> but it did not seem to work.</p>
<p>Thanks in advance.</p>

KFrank: <p>Hi Jazzy!</p>
<aside class="quote no-group" data-username="Jazzy66" data-post="1" data-topic="181345" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/jazzy66/48/61004_2.png" class="avatar"> Jazzy66:</div>
<blockquote>
<p>I find that the error is most likely caused by the last line of the code <code>V  =  V - 0.1 * (V_grad - VS_grad)</code></p>
</blockquote>
</aside>
<p>Yes, this is the cause of your problem.</p>
<p>There are two things going on here:</p>
<p>First, this line of code creates a <em>new</em> tensor (and then sets the name <code>V</code> to<br>
refer to it).  You almost certainly don’t want this.</p>
<p>Second, <code>V_grad</code> depends on (the original) <code>V</code> (as well as on <code>I</code> and <code>U</code>), and<br>
therefore on the computation graph that was built the first time through, but<br>
whose intermediate results have been freed, hence the “backward through<br>
the graph a second time” error.</p>
<p>You may solve the first problem by modifying <code>V</code> in place, and solve the second<br>
by wrapping the update of <code>V</code> in a <code>with torch.no_grad():</code> block.  Thus:</p>
<pre data-code-wrap="python"><code class="lang-python">    with torch.no_grad():
        V.sub_ (0.1 * (V_grad - VS_grad))
</code></pre>
<p>As an aside, although I don’t understand what you are trying to do, <code>VS =  V</code><br>
is probably a mistake, as <code>VS</code> and <code>V</code> are simply two names that reference the<br>
<em>same</em> tensor, so, for example, <code>VS_grad = gradient(I[i,:],U[i,:],VS)</code><br>
simply computes the same value as <code>V_grad</code>.</p>
<p>Also, at least in your sample code, setting <code>requires_grad = True</code> for <code>U</code> and<br>
<code>I</code> doesn’t do anything (except add some overhead) because you never use<br>
<code>U.grad</code> nor <code>I.grad</code>.</p>
<p>Best.</p>
<p>K. Frank</p>

Jazzy66: <p>Thank you, sir! you are my lifesaver and thanks again for your insightful and thorough answer!</p>
<p>I am sorry I did not mention that I also try to perform <a href="https://proceedings.neurips.cc/paper/2013/file/ac1dd209cbcc5e5d1c6e28598e8cbbe8-Paper.pdf" rel="noopener nofollow ugc">stochastic variation reduction gradient</a> in this loop, where <code>VS</code> is occasionally selected from <code>V</code> to compute the average gradient and regular gradient to reduce the stochastic variations across iterations. For the sake of simplicity, I did not put all these details into the fake code.</p>