Why is `retain_graph=True` needed in some case but not in other

magnusderrote: <p>Everywhere I read, people say just pass <code>retain_graph=True</code> to solve my issue, but I’d like to know what is under the hood. For example, the <code>loss.backward()</code> below doesn’t need to pass in any <code>retain_graph</code></p>
<pre><code class="lang-auto">class Mnist_Logistic(nn.Module):
    def __init__(self):
        super().__init__()
        self.lin = nn.Linear(784, 10)

    def forward(self, xb):
        return self.lin(xb)

for epoch in range(epochs):
    pred = model(input)
    loss = loss_func(pred, output)
    loss.backward()
</code></pre>
<p>For my model below, pytorch complains that it deletes the computational graph already, so if I run <code>backward()</code> in the next epoch, it could not access states variable. It seems that it can in the code above.</p>
<pre><code class="lang-auto">class Ode(nn.Module):
    def __init__(self, len_data, alpha=0.57, beta=0.11, ...):
        super().__init__()
        self.I, self.E, self.H = torch.zeros(1), torch.zeros(1), torch.zeros(1)
        self.S = nn.Parameter(torch.tensor([0.5]).to(device))
        self.sigma = nn.Parameter(torch.tensor([sigma])).to(device)
        # many nn.Parameter later

    def f(self, t, y):
        # only + and * parameters together

    def forward(self, I0, E0, H0):
        return torchdiffeq.odeint(self.f, t=time_range, y0=self.y0, method='rk4')


for data, y_exact in tqdm(train_dataloader):
    optimizer.zero_grad()    
    y_approx = model(data[0,0,0], data[0,0,1], data[0,0,2])
    loss = loss_fun(y_approx[1:, [1, 6, 7]], y_exact.squeeze())
    loss.backward()
    optimizer.step()
    scheduler.step()
</code></pre>

srishti-git1110: <p>Hi,<br>
PyTorch uses an aggressive memory freeing mechanism wherein as soon as a <code>.backward()</code> call is made on a tensor, PyTorch frees the “references to the saved tensors” in the computation graph associated with this tensor (on which .backward() is called).</p>
<p>These saved tensors are required to calculate the gradient of the tensor with respect to some other tensors, and hence a second backward call on the same tensor will produce an error as pytorch no longer has access to these saved tensors required for grad calculation.</p>
<p><code>retain_graph=True</code> causes pytorch not to free these references to the saved tensors.</p>
<p>So, in the first code that you posted, each time the <code>for</code> loop for training is run, a new computation graph is created - PyTorch uses dynamic graphs. This new graph saves references to tensors it’ll require for gradient computation. Hence, there’s no need to use retain_graph=True in this one.</p>
<p>As for the second code, please post the part where the .backward() calls are being made.</p>

magnusderrote: <p>Hi Srishti, thank you for your answer. I edited the question</p>

srishti-git1110: <p>Thanks for editing.<br>
But, could you please post a minimal executable snippet that would reproduce the error that you are facing.</p>
<p>From just skimming through this (second Model’s) code, I cannot see why this error would occur.</p>
<p>Anyway, were you able to understand the concept that I explained in my previous reply regarding what’s under the hood?</p>

ptrblck: <aside class="quote no-group" data-username="magnusderrote" data-post="1" data-topic="161705">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/m/5daacb/48.png" class="avatar"> magnusderrote:</div>
<blockquote>
<p>Everywhere I read, people say just pass <code>retain_graph=True</code> to solve my issue,</p>
</blockquote>
</aside>
<p>Which is unfortunately generally wrong.<br>
Do not pass <code>retain_graph=True</code> to any <code>backward</code> call unless you explicitly need it and can explain why it’s needed for your use case.<br>
Usually, it’s used as a workaround which will cause other issues afterwards. The mechanics of this argument were explained well by <a class="mention" href="/u/srishti-git1110">@srishti-git1110</a>.</p>

magnusderrote: <p>I managed to created an MRE like below. The <code>backward()</code> error is gone, but now I faced with an error that all the <code>param.grad</code> are 0, which I don’t know how to debug. Do you have any idea?</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
# this lib acts like scipy.integrate.ode but compatible with Torch tensors
from torchdiffeq.torchdiffeq import odeint   

class Ode(nn.Module):
    def __init__(self, beta=0.11, gamma=0.456):
        super().__init__()
        torch.set_default_dtype(torch.float64)
        self.I,  self.R = torch.zeros(1), torch.zeros(1)
        self.S = nn.Parameter(torch.tensor([0.5]))  # S is a hidden param
        self.y0 = torch.tensor([self.S, self.I, self.R])
        self.beta = nn.Parameter(torch.tensor([beta]))
        self.gamma = nn.Parameter(torch.tensor([gamma]))
        self.len_data = 2

    def f(self, t, y):
        S, I, R = y
        N = S + I + R
        return torch.tensor([S * I * self.beta / N,  # S
                             self.beta*I*S / N - self.gamma*I,   # I
                             self.gamma*I   # R
                             ], dtype=torch.float64, requires_grad=True)

    def forward(self, I0, R0):
        self.y0[1] = torch.Tensor([I0])
        self.y0[2] = torch.Tensor([R0])
        time_range = torch.linspace(0, self.len_data, self.len_data + 1)
        return odeint(self.f, t=time_range, y0=self.y0, method='rk4').double()


input = torch.DoubleTensor([[2,3], [4,5], [5,6]])
expected = torch.DoubleTensor([[[4, 5], [5, 6]], [[5, 6], [7, 8]], [[7, 8], [9, 10]]])
loss_fun = torch.nn.MSELoss()
model = Ode()
optimizer = torch.optim.Adam(model.parameters(), 1e-2)

for i in range(len(input)):
    optimizer.zero_grad()
    approx = model(input[i][0], input[i][1])
    observable_vars = approx[1:, [1, 2]]
    loss = loss_fun(observable_vars, expected[i])
    loss.backward()
    for param in optimizer.param_groups[0]['params']:
        if param.requires_grad:
            print(param.grad)
    optimizer.step()
</code></pre>

ptrblck: <p>Re-wrapping a tensor will detach the computation graph and create a new leaf tensor:</p>
<pre><code class="lang-python">        return torch.tensor([S * I * self.beta / N,  # S
                             self.beta*I*S / N - self.gamma*I,   # I
                             self.gamma*I   # R
                             ], dtype=torch.float64, requires_grad=True)
</code></pre>
<p>Are you seeing any issues if you are returning the result directly without re-creating a new tensor?</p>

magnusderrote: <p>The only issue is I didn’t know that creating a tensor will detach it from the computational graph <img src="https://discuss.pytorch.org/images/emoji/apple/see_no_evil.png?v=12" title=":see_no_evil:" class="emoji" alt=":see_no_evil:" loading="lazy" width="20" height="20">. Thank you</p>
<p>I changed that snippet to this</p>
<pre><code class="lang-auto">return torch.cat([S * I * self.beta / N,  # S
                             self.beta*I*S / N - self.gamma*I,   # I
                             self.gamma*I   # R
                         ])
</code></pre>
<p>Now</p>
<pre><code class="lang-auto">    optimizer.zero_grad()
    ...
    for name, param in model.named_parameters():
        print(name, param.requires_grad, param.grad)
    print(loss)
    ...
    optimizer.step()
</code></pre>
<p>will print</p>
<pre><code class="lang-auto">S True None
beta True tensor([-0.5961])
gamma True tensor([2.9083])
tensor(7.4195, grad_fn=&lt;MseLossBackward0&gt;)
S True None
beta True tensor([-0.6962])
gamma True tensor([11.5022])
tensor(8.7388, grad_fn=&lt;MseLossBackward0&gt;)
S True None
beta True tensor([-1.0248])
gamma True tensor([17.9899])
tensor(15.6819, grad_fn=&lt;MseLossBackward0&gt;)
</code></pre>
<p>I still don’t get why <code>S.grad</code> is <code>None</code></p>

magnusderrote: <p>Thank you for the explanation. Yes, I was able to understand the mechanism of <code>retain_graph=True</code>. What you said, and I didn’t know, is that</p>
<blockquote>
<p>So, in the first code that you posted, each time the <code>for</code> loop for training is run, a new computation graph is created - PyTorch uses dynamic graphs.</p>
</blockquote>
<p>Does that mean that since the first code use some pytorch layers, they got that recreation logic baked in?</p>

srishti-git1110: <aside class="quote no-group" data-username="magnusderrote" data-post="10" data-topic="161705">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/m/5daacb/48.png" class="avatar"> magnusderrote:</div>
<blockquote>
<p>Does that mean that since the first code use some pytorch layers, they got that recreation logic baked in?</p>
</blockquote>
</aside>
<p>It only means that each time an iteration of the training loop is run, a fresh graph is created for the loss tensor (i.e. after the backward call on it, the previous graph’s buffers are freed) - this is called dynamic graph creation, in the sense that graphs get created from scratch and on the go each time.</p>
<p>And since, a new graph is getting created after every backward call, there’s no need to retain the previous graph.</p>

srishti-git1110: <p>I haven’t used odeint ever, but from inspecting your code, I think the issue might be the <code>loss</code> tensor not being a function of <code>self.S</code>, which means <code>self.S</code> might not be getting used anywhere in the calculation of loss.</p>
<p>This could be a possible cause for <code>self.S</code> 's grad attribute being None after the backward call on <code>loss</code>.</p>
<p>The reason why I think so is that I can see <code>self.beta</code> and <code>self.gamma</code> being used in the method <code>f</code> to do some calculations.</p>
<p>But, I cannot see <code>self.S</code> being used anywhere that might be causing the final tensor loss to be free from <code>self.S</code>.</p>
<p>Not very sure though. I might be wrong.</p>

magnusderrote: <p>Thank you both. I used your hints and fixed the problem. Now all parameters have gradients flow through them. Since I can only set one solution, I don’t know which post to set, please let me know <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=12" title=":smiley:" class="emoji" alt=":smiley:" loading="lazy" width="20" height="20">. Right now I am just giving all of them a heart.</p>

LLlearner: <aside class="quote no-group" data-username="srishti-git1110" data-post="2" data-topic="161705">
<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/srishti-git1110/48/52515_2.png" class="avatar"> srishti-git1110:</div>
<blockquote>
<p>, a new computation graph is created - PyTorch uses dynamic graphs</p>
</blockquote>
</aside>
<p>You said that for the first case a new computation graph is created - PyTorch uses dynamic graphs.  The backward pass from loss involves the parameters of the network W, which are trainable. If PyTorch frees the “references to the saved tensors”, means it frees the reference to W, as well. How are the weights incrementally update then? And not initialized again for every step?</p>

srishti-git1110: <p>Hi, <a class="mention" href="/u/lllearner">@LLlearner</a></p>
<aside class="quote no-group" data-username="LLlearner" data-post="14" data-topic="161705">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/l/cab0a1/48.png" class="avatar"> LLlearner:</div>
<blockquote>
<p>If PyTorch frees the “references to the saved tensors”, means it frees the reference to W, as well.</p>
</blockquote>
</aside>
<p>The weights W of any network are leaf tensors in the computation graph of the calculated loss.<br>
As soon as <code>loss.backward()</code> is called the gradient calculation happens by backprop and the <code>grad</code> attribute of the leaf tensors W is populated by default, and the references to the saved tensors (saved tensors are the ones that’ll be required to compute the gradients of the loss wrt the weights W) are freed, and so the gradient cannot be computed on a second backward call thus giving an error.</p>
<aside class="quote no-group" data-username="LLlearner" data-post="14" data-topic="161705">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/l/cab0a1/48.png" class="avatar"> LLlearner:</div>
<blockquote>
<p>How are the weights incrementally update then?</p>
</blockquote>
</aside>
<p>The (optimisable) weights W are then updated using <code>optimizer.step()</code> making use of the <code>grad</code> attribute.</p>
<aside class="quote no-group" data-username="LLlearner" data-post="14" data-topic="161705">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/l/cab0a1/48.png" class="avatar"> LLlearner:</div>
<blockquote>
<p>And not initialized again for every step?</p>
</blockquote>
</aside>
<p>On the next iteration of the training loop, a new graph is formed <em>using the updated weights W</em> (hence no re-initialization) and this continues.</p>
<p>Hope this helps,<br>
S</p>