How to manually update network parameters while keeping track of its computational graph?

gunnxx: <p>I am trying to implement MAML. With simplification, these are the three operations I wish to implement where <code>theta</code> is the weights of a neural network.<br>
<a href="https://www.codecogs.com/eqnedit.php?latex=%5Cdpi%7B100%7D&amp;space;%5Csmall&amp;space;%5Cbegin%7B%7D%5C%5C&amp;space;%5Ctheta_1&amp;space;%5Cleftarrow&amp;space;%5Ctheta&amp;space;-&amp;space;%5Calpha&amp;space;%5Cfrac%7B%5Cpartial&amp;space;f(%5Ctheta)%7D&amp;space;%7B%5Cpartial&amp;space;%5Ctheta%7D%5C%5C&amp;space;%5Ctheta_2&amp;space;%5Cleftarrow&amp;space;%5Ctheta_1&amp;space;-&amp;space;%5Calpha&amp;space;%5Cfrac%7B%5Cpartial&amp;space;f(%5Ctheta_1)%7D%7B%5Cpartial&amp;space;%5Ctheta_1%7D%5C%5C&amp;space;%5Ctheta&amp;space;%5Cleftarrow&amp;space;%5Ctheta&amp;space;-&amp;space;%5Calpha&amp;space;%5Cfrac%7B%5Cpartial&amp;space;f(%5Ctheta_2)%7D%7B%5Cpartial&amp;space;%5Ctheta%7D&amp;space;%5Cend%7B%7D" target="_blank" rel="noopener nofollow ugc"><img src="https://latex.codecogs.com/gif.latex?%5Cdpi%7B100%7D&amp;space;%5Csmall&amp;space;%5Cbegin%7B%7D%5C%5C&amp;space;%5Ctheta_1&amp;space;%5Cleftarrow&amp;space;%5Ctheta&amp;space;-&amp;space;%5Calpha&amp;space;%5Cfrac%7B%5Cpartial&amp;space;f(%5Ctheta)%7D&amp;space;%7B%5Cpartial&amp;space;%5Ctheta%7D%5C%5C&amp;space;%5Ctheta_2&amp;space;%5Cleftarrow&amp;space;%5Ctheta_1&amp;space;-&amp;space;%5Calpha&amp;space;%5Cfrac%7B%5Cpartial&amp;space;f(%5Ctheta_1)%7D%7B%5Cpartial&amp;space;%5Ctheta_1%7D%5C%5C&amp;space;%5Ctheta&amp;space;%5Cleftarrow&amp;space;%5Ctheta&amp;space;-&amp;space;%5Calpha&amp;space;%5Cfrac%7B%5Cpartial&amp;space;f(%5Ctheta_2)%7D%7B%5Cpartial&amp;space;%5Ctheta%7D&amp;space;%5Cend%7B%7D" title="\small \begin{}\\ \theta_1 \leftarrow \theta - \alpha \frac{\partial f(\theta)} {\partial \theta}\\ \theta_2 \leftarrow \theta_1 - \alpha \frac{\partial f(\theta_1)}{\partial \theta_1}\\ \theta \leftarrow \theta - \alpha \frac{\partial f(\theta_2)}{\partial \theta} \end{}" width="116" height="65"></a><br>
Assumes constant input i.e. <code>f(theta) = loss(net(in_))</code> because we only interested in the gradients w.r.t. weights.</p>
<p>Here’s my code snippet</p>
<pre><code class="lang-python">temp_net = copy.deepcopy(net)

## First formula
loss_1 = loss(net(in_))
grad_1 = torch.autograd.grad(loss_1, net.parameters(), create_graph=True)

with torch.no_grad():
    for param, grad in zip(temp_net.parameters(), grad_1):
        new_param = param - lr * grad
        param.copy_(new_param)

## Second formula
loss_2 = loss(temp_net(in_))
grad_2 = torch.autograd.grad(loss_2, temp_net.parameters(), create_graph=True)

with torch.no_grad():
    for param, grad in zip(temp_net.parameters(), grad_2):
        new_param = param - lr * grad
        param.copy_(new_param)

## Third formula
loss_3 = loss(temp_net(in_))
grad_3 = torch.autograd.grad(loss_3, net.parameters())
</code></pre>
<p>Last line throws an error because the computational graph disconnects when I use <code>with torch.no_grad()</code>. I can’t compute gradient w.r.t. <code>net.parameters()</code>. However, if I remove the <code>torch.no_grad()</code>, it throws an error <code>RuntimeError: a leaf Variable that requires grad is being used in an in-place operation.</code></p>
<p>My question is how can I implement the given formula? How can I update my network parameters manually while maintaining its computational graph?</p>

gunnxx: <p>Seeing from <code>learn2learn</code> implementation of MAML, turns out we can manually do in-place update of the parameters (not using <code>.copy_()</code>) to maintain the computational graph. However, it requires us to recursively check the <code>._parameters</code> property of each <code>_modules</code> of our model. This is <a href="https://github.com/learnables/learn2learn/blob/06893e847693a0227d5f35a6e065e6161bb08201/learn2learn/utils/__init__.py#L283" rel="noopener nofollow ugc">the manual in-place update</a> and this is <a href="https://github.com/learnables/learn2learn/blob/06893e847693a0227d5f35a6e065e6161bb08201/learn2learn/utils/__init__.py#L297" rel="noopener nofollow ugc">the recursive part</a>.</p>