Example of second order derivative (single variable)

username: Luvata
cooked: <p>I’m implementing MAML, but currently I’m struggling to understand how to calculate the gradient of the outer loss w.r.t model’s parameters before any inner gradient step. So I made a naive example, but I’m not sure if I missed anything ?</p>
<pre><code class="lang-auto">import torch
from torch.autograd import grad

# f(x) = 2x + 3

x_train = torch.tensor(2., requires_grad=True)
y_train = torch.tensor(7., requires_grad=True)  # 2 * 2 + 3 = 7

x_test = torch.tensor(0.2, requires_grad=True)
y_test  = torch.tensor(3.4, requires_grad=True) # 0.2*2 + 3 = 3.4

W = torch.tensor(1., requires_grad=True)

n_inner_update = 3

for step in range(2000):
    y_pred = x_train * W + 3.
    loss = torch.abs(y_pred - y_train)
    inner_grad = grad(loss, W, create_graph=True)[0]

    W_after_update = W - inner_grad * 0.002  # an inner update

    for inner in range(n_inner_update - 1):
        y_pred = x_train * W_after_update + 3.
        loss = torch.abs(y_pred - y_train)
        inner_grad = grad(loss, W_after_update, create_graph=True)[0]
        W_after_update = W_after_update - inner_grad* 0.002

    y_pred_test = x_test * W_after_update + 3.
    outer_loss = torch.abs(y_pred_test - y_test)
    outer_loss_wrt_W = grad(outer_loss, W)[0]
    W.data -= outer_loss_wrt_W * 0.01
    print(loss.item())
</code></pre>

username: albanD
cooked: <p>Note that the right way to update the weight is with:</p>
<pre><code class="lang-auto">with torch.no_grad():
    W -= outer_loss_wrt_W * 0.01
</code></pre>
<p><code>.data</code> should not be used anymore!</p>
<p>Your code looks good at first glance. What is the issue that you have with it?</p>

username: Luvata
cooked: <p>Thank you for your answer, I knew that I missed something!<br>
This is my toy example when trying to under stand MAML and I think I kind of understand and familiar with pytorch.<br>
When implement MAML, I have a model (nn.Module) with parameters W (stack of CNN and Linear). How can I still <strong>keep my W</strong> but also <strong>forward with <code>W_after_update</code></strong> so that I can still backward the outer_loss to my starting W ?</p>

username: albanD
cooked: <p>I’m afraid it is a tricky thing to do.<br>
If you follow MAML quite closely, I would recommend using <a href="https://github.com/facebookresearch/higher/">https://github.com/facebookresearch/higher/</a> that takes care of all this state management for you.</p>

username: Luvata
cooked: <p>Thank you for your reply ! Higher looks really promising<br>
Since I understood the concept of pytorch grad, I will try to implement it with normal pytorch (I think about deep copy and passing back grad_output)<br>
Once again thank you so much for your help</p>

username: albanD
cooked: <p>You can do it with pytorch (that’s what higher does <img src="https://discuss.pytorch.org/images/emoji/apple/stuck_out_tongue.png?v=9" title=":stuck_out_tongue:" class="emoji" alt=":stuck_out_tongue:"> ).<br>
Be careful that the nn module is not built for that unfortunately.<br>
keep in mind in particular:</p>
<ul>
<li>nn.Parameter can only contains leafs (gradient with no history).</li>
<li>optim.step() is NOT differentiable and won’t propagate gradients.</li>
</ul>
<p>Good luck!<br>
Don’t hesitate to post here if you need inputs.</p>