Reuse the same tensor causes RuntimeError: Trying to backward through the graph a second time (or directly access saved tensors after they have already been freed)

username: Yuchen_Shi
cooked: <p>Here’s the example code:</p>
<pre><code class="lang-python">import torch
import torch.nn as nn
class Model(nn.Module):
    def __init__(self, d_model):
        super(Model, self).__init__()
        self.linear1 = nn.Linear(d_model, d_model)
        self.linear2 = nn.Linear(d_model, 1)
    def forward(self, x):
        return self.linear1(x)
    def calculate(self, x):
        return self.linear2(x)
batch_size, d_model, target = 4, 3, 1
data = torch.rand(batch_size, d_model)
model = Model(d_model)
temp = model(data)  # (d_model, d_model)
optimizer = torch.optim.Adam(model.parameters())
for i in range(10):
    y = model.calculate(temp)  # (d_model, 1)
    loss = ((y - target) ** 2).mean()
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
</code></pre>
<p>Assume that the <code>forward</code> of the model is so complex that it may not be computed in each for loop. I get the error <code>RuntimeError: Trying to backward through the graph a second time (or directly access saved tensors after they have already been freed). Saved intermediate values of the graph are freed when you call .backward() or autograd.grad(). Specify retain_graph=True if you need to backward through the graph a second time or if you need to access saved tensors after calling backward.</code></p>
<p>I think <code>backward</code> destroy the graph of <code>forward</code> so that it cannot be called in the second and subsequent updates. <code>loss.backward(retain_graph=False)</code> solves the problem, but is there other efficient way?  Any help would be much appreciated.</p>

username: ptrblck
cooked: <p>The intermediate tensors created in <code>forward</code> are freed in the <code>backward</code> call and the next iteration will fail since you are only recomputing the <code>calculate</code> part of your model.<br>
Could you explain your use case a bit more and why <code>forward</code> is not recomputed anymore?</p>

username: Yuchen_Shi
cooked: <p>Thanks for your reply!</p>
<p>I’m applying the <a href="https://arxiv.org/abs/1707.06347" rel="noopener nofollow ugc">Proximal Policy Optimization Algorithms</a> to my problem following the code <a href="https://github.com/vwxyzjn/ppo-implementation-details/blob/main/ppo.py" rel="noopener nofollow ugc">vwxyzjn/ppo-implementation-details</a>. The alrorithm is shown below:</p>
<p><img src="https://i.stack.imgur.com/a6z3u.png" alt="" role="presentation" width="690" height="195"></p>
<p>where the advantage estimates A is</p>
<ol>
<li>used to obtain surrogate loss function L;</li>
<li>calculated by current value function V parameterized by φ;</li>
<li>repeatedly used in K epochs.</li>
</ol>
<p>The input of V is usually the observation of an agent (flattened data without gradients), but in my case the input of V is the output of a parameterized deep network used to aggregate features. I have modified some naming of the code to make it clearer.</p>
<pre><code class="lang-python">import torch
import torch.nn as nn
class Model(nn.Module):
    def __init__(self, d_model):
        super(Model, self).__init__()
        self.linear1 = nn.Linear(d_model, d_model)
        self.linear2 = nn.Linear(d_model, 1)
    def feature_aggregation(self, x):
        return self.linear1(x)
    def value_function(self, x):
        return self.linear2(x)
batch_size, d_model, target = 4, 3, 1
observation = torch.rand(batch_size, d_model)
model = Model(d_model)
h_observation = model.feature_aggregation(observation)  # (d_model, d_model)
optimizer = torch.optim.Adam(model.parameters())
for i in range(10):
    value = model.value_function(h_observation)  # (d_model, 1)
    loss = ((value - target) ** 2).mean()
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
</code></pre>
<p>The <code>feature_aggregation</code> (i.e. <code>forward</code> in the original code) is calculated once because I want to save running time. Please let me explain in detail my thoughts, but forgive me for my lack of understanding of the <code>backward</code> mechanism:</p>
<ul>
<li>since the input of <code>feature_aggregation</code> is the same for all epochs, I do not want to spend repeated time calculating it, but instead to keep the computational graph that leads to <code>h_observation</code>;</li>
<li>whenever <code>value</code> is calculated, I want to calculate the computational graph for <code>value</code> “on the basis of” the previously saved computational graph for <code>h_observation</code>;</li>
<li>during <code>backward</code>, I wonder if I can free the computational graph of the <code>value_function</code> part because it is useless and this may save time and memory.</li>
</ul>

username: ptrblck
cooked: <p>Based on your explanation it seems <code>h_observation</code> is treated as a static input and you don’t want to optimizer the <code>feature_aggregation</code> step in each iteration. If so, you could execute this step in a <code>no_grad</code> guard to avoid creating the computation graph:</p>
<pre><code class="lang-python">with torch.no_grad():
    h_observation = model.feature_aggregation(observation)
</code></pre>
<p>which will create <code>h_observation</code> as a tensor without any Autograd history.</p>