Backward fails when I use a composition for torch.nn.Parameter

n_kotel: <p>Hello,<br>
Initially, I faced the following issue:</p>
<p><em>RuntimeError: Trying to backward through the graph a second time, but the buffers have already been freed. Specify retain_graph=True when calling backward the first time.</em></p>
<p>I figured out that the problem raised when I use torch.nn.Parameter in my model and do some transformations with it. It is easy to understand just by looking at the toy example:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn

# Define our 'model'
class Simple_func(nn.Module):
    def __init__(self,):
        super(Simple_func, self).__init__()
        self.a = nn.Parameter((torch.randn(1) + 5) ** 2)
#         self.b = torch.log(self.a)
    def forward(self, x):
        return x + self.a
#         return x + self.b

model = Simple_func()
print([p for p in model.parameters()])
# here I have the only parameter, it's value is (torch.randn(1) + 5) ** 2

# And here we have 'training process.'

optim = torch.optim.Adamax(model.parameters())
data = torch.arange(10, dtype=torch.float32) - 5
for i, d_point in enumerate(data):
    loss = (model(d_point) - d_point)**2
    loss.backward()
    optim.step()
    optim.zero_grad()

</code></pre>
<p>The code above works fine.</p>
<p>But let’s say for numerical stability I want to optimise such a Parameter which will be inside a logarithm function, so I change only Simple_func in the code above:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn

# Define our 'model'
class Simple_func(nn.Module):
    def __init__(self,):
        super(Simple_func, self).__init__()
        self.a = nn.Parameter((torch.randn(1) + 5) ** 2)
        self.b = torch.log(self.a) . # &lt;-------- UNCOMMENTED
    def forward(self, x):
#         return x + self.a # &lt;-------- COMMENTED
        return x + self.b # &lt;-------- UNCOMMENTED

model = Simple_func()
print([p for p in model.parameters()])
# here I have the only parameter, it's value is (torch.randn(1) + 5) ** 2

# And here we have 'training process'

optim = torch.optim.Adamax(model.parameters())
data = torch.arange(10, dtype=torch.float32) - 5
for i, d_point in enumerate(data):
    loss = (model(d_point) - d_point)**2
    loss.backward()
    optim.step()
    optim.zero_grad()

</code></pre>
<p>The code above does not work anymore. It throws the error at the second backward:</p>
<p><em>RuntimeError: Trying to backward through the graph a second time, but the buffers have already been freed. Specify retain_graph=True when calling backward the first time.</em></p>
<p>My questions are:</p>
<ol>
<li>What is the problem which causes such an error?</li>
<li>What is the correct way to use nn.Parameter to encapsulate it in other functions?</li>
</ol>
<p>I understand that I could use torch.tensor notation here with requires_grad=True, but as far as I understand I will have to manually add such tensors to torch.optim to optimise over them.</p>

tom: <aside class="quote no-group" data-username="n_kotel" data-post="1" data-topic="58755">
<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/n/bcef8e/48.png" class="avatar"> n_kotel:</div>
<blockquote>
<ol>
<li>What is the problem which causes such an error?</li>
<li>What is the correct way to use nn.Parameter to encapsulate it in other functions?</li>
</ol>
</blockquote>
</aside>
<ol>
<li>The problem is that you only compute self.b once - in the <code>__init__</code>. Then on the first backward, it’ll work. But the next time around, the edge in the computation graph connecting <code>self.b</code> to <code>self.a</code> will be gone.</li>
<li>The most obvious way to get this to work is to move the line <code>self.b = torch.log(self.a)</code> from the <code>__init__</code> to the top of the forwrd.</li>
<li>This comes up every <a href="https://github.com/pytorch/pytorch/issues/7313" rel="noopener nofollow ugc">once in a while</a>, but it’s not important enough/the workaround is good enough so it is not a priority.</li>
</ol>
<p>Best regards</p>
<p>Thomas</p>

n_kotel: <p>I got it, thank you!</p>