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 ?
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())