Cannot calculate second order gradients even though `create_graph=True`

username: jwillette
cooked: <p>I have a training loop iteration that looks like this…</p>
<pre><code class="lang-auto">mu, logvar = m(x)

alpha = torch.zeros(size, requires_grad=True)
loss = alpha * criterion(mu, logvar)

loss.backward(retain_graph=True, create_graph=True)
for p in m.parameters():
  p = p - LR * p.grad

x_next, y_next = data[j + 1]
mu_next, logvar_next = m(x_next)

meta_loss = criterion(mu_next, logvar_next)
grad_alpha = torch.autograd.grad(meta_loss, alpha)[0]
</code></pre>
<p>and pytorch is complaining that <code>RuntimeError: One of the differentiated Tensors appears to not have been used in the graph. Set allow_unused=True if this is the desired behavior.</code> but as far as I can tell it should have been used in the graph because it was multiplied by the loss and then the graph should have persisted, and then when calling the model again I should be able to get the gradients with respect to the original <code>alpha</code>.</p>
<p>Where have I gone wrong?</p>

username: ptrblck
cooked: <p>The computation graph, where <code>alpha</code> was used seems to be unrelated to the second computation graph, as you are using new input tensors and I cannot see the connection between these graphs.</p>

username: jwillette
cooked: <p>Shouldn’t the connection be the model parameters themselves?</p>

username: ptrblck
cooked: <p>If <code>alpha</code> was created through the previous graph, then yes.<br>
Otherwise, if <code>alpha</code> is a leaf node, then it won’t connect the computation graphs.</p>

username: jwillette
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> thanks for the responses. Maybe I am just dense, but I am not understanding what you said. I’ve been through this stuff many times in the docs but I still don’t get exactly why this isn’t working. I have updated my example to be more simple and complete.</p>
<p>Does my updated example look like it would work? Why or why not? Thanks again</p>

username: ptrblck
cooked: <p>I’ve added some annotation to your code and added a dummy model:</p>
<pre><code class="lang-python">class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.fc1 = nn.Linear(1, 1)
        self.fc2 = nn.Linear(1, 1)
        
    def forward(self, x):
        mu = self.fc1(x)
        logvar = self.fc2(x)
        return mu, logvar


# Setup
m = MyModel()
criterion = nn.MSELoss()
x = torch.randn(1, 1)
LR = 1.

# First graph creation
mu, logvar = m(x)

alpha = torch.zeros(size, requires_grad=True)
print(alpha.is_leaf) # alpha is a leaf variable
loss = alpha * criterion(mu, logvar) # alpha is now used in the computation graph via the multiplication

loss.backward(retain_graph=True, create_graph=True)
print(alpha.grad) # alpha has a gradient
for p in m.parameters():
  p = p - LR * p.grad

# Next graph
x_next = torch.randn(1, 1)
mu_next, logvar_next = m(x_next) # this creates a new computation graph
# no outputs from the first computation were used here

meta_loss = criterion(mu_next, logvar_next) # loss calculated for the new graph
# calculate sum of grads of meta_loss w.r.t. alpha,
# but alpha wasn't used in the creation of meta loss!
grad_alpha = torch.autograd.grad(meta_loss, alpha)[0] 
</code></pre>
<p>The main issue is, that you are trying to calculate the gradients of <code>meta_loss</code> w.r.t. <code>alpha</code>, but when you write down the formula, how <code>meta_loss</code> was computed, you won’t find <code>alpha</code> in there.</p>

username: jwillette
cooked: <p>ok, that is very clear now…but… alpha was used in the calculation of the parameters of the model before <code>m(x_next)</code> was called. I’m assuming that the graph containing this information was either never created  or destroyed?</p>

username: ptrblck
cooked: <p>The graph was created and <code>loss</code> is attached to it.<br>
Using <code>torchviz</code> you would get this graph for the first forward pass with <code>alpha</code> included:</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/d/9/d9fb425f0ccd91f3fe4bfe63995ad2bbdd72227d.png" alt="graph_with_alpha" data-base62-sha1="v6lVKYuXnr1LP4l71p8ZMHVCUtn" width="333" height="410"></p>
<p>and this graph for the second part:</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/e/4/e42d7c0fc2a31e43eae8f840e2669c0cf555cd12.png" alt="graph2_wo_alpha" data-base62-sha1="wyyjru2UHarYT3E3nQ2tWXFw17I" width="333" height="340"></p>
<p>As you can see, <code>alpha</code> is in the graph in the last multiplication.<br>
However, these graphs are not depending on each other, as they were created in two separate forward passes and new input.</p>

username: jwillette
cooked: <p>right. that is a very helpful visualization. I see what is happening and it is very clear to me that the gradients don’t make it past the backward call. I still don’t see <em>why</em> they dont propagate through the parameter update though.</p>
<ul>
<li>The parameters require grad</li>
<li>The loss required grad</li>
<li>I set <code>retain_graph</code> and <code>create_graph</code>
</li>
<li>I manually update the parameters (without calling <code>no_grad</code>) and I also don’t use the in-place operations that are in the optimizers which avoid the differentiation graph (AFAIK)</li>
</ul>
<p>This chain of events leads me to believe that the graph should still be stored and the top graph you posted and the bottom graph should be connected from <code>MeanBackward0</code> to <code>TBackward</code> with the inputs coming in separately.</p>
<p>This is a common pattern in meta learning if I understand it correctly and I don’t see <em>why</em> the graphs should be separated given my bullet points above.</p>

username: ptrblck
cooked: <p>I think your bullet points are correct.<br>
Also,</p>
<aside class="quote no-group" data-username="jwillette" data-post="9" data-topic="78711">
<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/jwillette/48/17366_2.png" class="avatar"> jwillette:</div>
<blockquote>
<p>This chain of events leads me to believe that the graph should still be stored</p>
</blockquote>
</aside>
<p>would be true.</p>
<p>However, <code>MeanBackward0</code> and <code>TBackward</code> will not connect the graphs.</p>
<p>If these graphs would be connected, then also this dummy code:</p>
<pre data-code-wrap="python"><code class="lang-python">lin = nn.Linear(1, 1)

x1 = torch.randn(1, 1)
out1 = lin(x1)

x2 = torch.randn(1, 1)
out2 = lin(x2)
</code></pre>
<p>would create a single graph, which is not the case.</p>
<aside class="quote no-group" data-username="jwillette" data-post="9" data-topic="78711">
<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/jwillette/48/17366_2.png" class="avatar"> jwillette:</div>
<blockquote>
<p>This is a common pattern in meta learning if I understand it correctly and I don’t see <em>why</em> the graphs should be separated given my bullet points above.</p>
</blockquote>
</aside>
<p>I have the feeling, that I’m missing something, as I’m not deeply familiar with meta learning.</p>
<p>CC <a class="mention" href="/u/alband">@albanD</a> to see what I’m missing and correct me.</p>

username: jwillette
cooked: <p>The method in the following two papers, as well as the corresponding github link use second order gradients in this way.</p>
<ul>
<li><a href="https://arxiv.org/abs/1703.03400" rel="nofollow noopener">https://arxiv.org/abs/1703.03400</a></li>
<li><a href="https://arxiv.org/abs/1803.09050" rel="nofollow noopener">https://arxiv.org/abs/1803.09050</a></li>
</ul>
<p>I think the torchmeta library in this link is doing something in the background to allow this to happen, but I can’t see where the disconnect is and why my own code never seems to recognize the second order gradients when I think it should</p>
<ul>
<li><a href="https://github.com/tristandeleu/pytorch-meta/blob/master/examples/maml/train.py#L67" rel="nofollow noopener">https://github.com/tristandeleu/pytorch-meta/blob/master/examples/maml/train.py#L67</a></li>
</ul>

username: albanD
cooked: <p>I think the problem in your code is:</p>
<pre><code class="lang-auto">for p in m.parameters():
  p = p - LR * p.grad
</code></pre>
<p>This does not modify <code>p</code> inplace ! It just assigns the result to a variable name p that you override just after. So you could remove these lines and your code would run the same.<br>
This is why you don’t see the link.</p>
<p>Another issue you’re gonna face is that <code>nn.Parameter()</code> are built explicitly to be leaf Tensors (with no history) and so you won’t be able to change them with a Tensor you try to backprop through.</p>
<p>You can check the <a href="https://github.com/facebookresearch/higher">higher</a> package to do this properly (they solve this problem of nn.Parameter for you). Otherwise, if you want to do it manually, you can find some info in this thread: <a href="https://discuss.pytorch.org/t/how-does-one-have-the-parameters-of-a-model-not-be-leafs/70076/9" class="inline-onebox">How does one have the parameters of a model NOT BE LEAFS?</a></p>

username: jwillette
cooked: <p><a class="mention" href="/u/alband">@albanD</a> <a class="mention" href="/u/ptrblck">@ptrblck</a> thanks for taking the time to clarify this</p>