Gradient computation in meta-learning algorithms

rinuboney: <p>Hi,</p>
<p>Related question: <a href="//discuss.pytorch.org/t/resolved-implementing-maml-in-pytorch/4053">http://discuss.pytorch.org/t/resolved-implementing-maml-in-pytorch/4053</a></p>
<p>I’m trying to implement meta-learning algorithms such as <a href="https://arxiv.org/abs/1703.03400" rel="nofollow noopener">MAML</a>(<a href="https://github.com/cbfinn/maml/blob/master/maml.py" rel="nofollow noopener">TF implementation</a>) and <a href="https://arxiv.org/pdf/1707.09835.pdf" rel="nofollow noopener">Meta-SGD</a> in PyTorch. I have difficulties in understanding the gradient flow.</p>
<pre><code class="lang-auto">train_samples, test_samples = task['train'], task['test']

images, labels = train_samples['image'], train_samples['label']
if args.cuda:
    images, labels = images.cuda(), labels.cuda()
images, labels = Variable(images), Variable(labels)

images, labels = test_samples['image'], test_samples['label']
if args.cuda:
    images, labels = images.cuda(), labels.cuda()
images, labels = Variable(images), Variable(labels)

# Inner gradient update

alpha = [Variable(p.data.clone().fill_(0.01), requires_grad=True) for p in base_learner.parameters()]

meta_optimizer = optim.Adam(base_learner.parameters(), lr=args.update_lr)
meta_optimizer.zero_grad()

output = base_learner(images)
loss_fn = nn.CrossEntropyLoss()
loss = loss_fn(output, labels)
loss.backward(retain_graph=True)

for param in base_learner.parameters():
    param.grad.requires_grad = True

base_learner_ = BaseLearner(args.num_classes, args.num_filters)
if args.cuda:
    base_learner_.cuda()
base_learner_.train()

print(list(base_learner_.parameters())[0].data.sum())

for param_, lr, param in zip(base_learner_.parameters(), alpha, base_learner.parameters()):
    param_ = param - lr*param.grad

print(list(base_learner_.parameters())[0].data.sum())

# Note: parameters of base_learner_ doesn't seem to get updated here.

# Meta update

base_learner_.zero_grad()

output = base_learner_(images)
loss_fn = nn.CrossEntropyLoss()
loss = loss_fn(output, labels)
loss.backward()

meta_optimizer.step()
</code></pre>
<ol>
<li>When I set the parameters of <code>base_learner_</code> from <code>base_learner</code>, that doesn’t seem to be executed. How to do this properly?</li>
<li>The gradient while backpropagating from <code>base_learner_</code> doesn’t flow to the parameters of <code>base_learner</code>. I know that this is because the parameters of <code>base_learner_</code> are not set properly based on <code>base_learner</code>. I want to set it such that gradients of <code>base_learner_</code> flows to parameters of <code>base_learner</code>.</li>
</ol>
<p>The algorithm I’m trying to implement can be found at the beginning of page 5 here - <a href="https://arxiv.org/pdf/1707.09835.pdf" rel="nofollow noopener">Meta-SGD: Learning to Learn Quickly for Few Shot Learning</a>. Is this the right approach? What is the pragmatic way of doing it?</p>
<p>More code:</p>
<pre><code class="lang-auto">
# Model

class ConvModule(nn.Module):
    """ Conv Module """
    def __init__(self, in_channels, out_channels):
        super(ConvModule, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
        self.bn = nn.BatchNorm2d(out_channels)

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = F.relu(x)
        return F.max_pool2d(x, 2)

class BaseLearner(nn.Module):
    """ Simple Conv Net """
    def __init__(self, num_classes, num_filters=64):
        super(BaseLearner, self).__init__()
        self.conv1 = ConvModule(1, num_filters)
        self.conv2 = ConvModule(num_filters, num_filters)
        self.conv3 = ConvModule(num_filters, num_filters)
        self.conv4 = ConvModule(num_filters, num_filters)
        self.fc = nn.Linear(num_filters, num_classes)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        x = self.conv4(x)
        x = x.view(-1, 64)
        return self.fc(x)
</code></pre>

aitutakiv: <p>I’m not sure why you can’t manually set the parameters of base_learner_<br>
But regardless, for MAML one problem you have is that in the meta step, you need to compute the loss using the updated weights, but then you need to compute the gradients with respect to the <em>original</em> weights.<br>
To achieve this I just did away with modules, and directly used the functions that allow you to pass in different weight variables as arguments.<br>
I also found this function useful for collecting gradients for the meta step:<br>
<a href="http://pytorch.org/docs/master/_modules/torch/autograd.html#grad" class="onebox" target="_blank" rel="nofollow noopener">http://pytorch.org/docs/master/_modules/torch/autograd.html#grad</a></p>

rinuboney: <p>Yes I can replace the Conv2d module with the functional form but then how did you deal with the affine parameters in batch norm?</p>

dragen: <p>Hello, Have you implemented the algorithm of Meta-SGD successfully yet?  we can communicate since I’m working on Meta-SGD also.</p>

CodePothunter: <p>yep, It will be useful with <code>functional</code> if it’s a conv-like net, but what should I do when I want to apply a meta update in a recurrent network?</p>

Sudarshan_VB: <p><a class="mention" href="/u/aitutakiv">@aitutakiv</a>, I understand how doing away with modules makes it easier, but writing a functional forward for deeper networks is a pain, so is there a way where to code maml or foml using modules ?</p>
<p>Please let me know, any help is appreciated !</p>

Miao_txy_Aaron: <p>Hi, have you implemented the MAML in “module” way. Hope to post a link, thks</p>

asura: <p>Hello <a class="mention" href="/u/aitutakiv">@aitutakiv</a>, I am a bit confused. How should I compute the gradients wrt to the original weights? If it’s possible, can you review my code. It seems I am approaching the problem quite naively, and missing out on something important! It would be great if you are able to give me a few pointers <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>
<p>My inner loop code:</p>
<pre><code class="lang-auto">def inner_loop(learner, batch, loss_fn, steps, device):
    X, y = batch # Batch Size = 2*shots*ways
    X, y = X.to(device), y.to(device)

    # Separate data into adaptation/evalutation sets
    # Dtrain and Dtest both of size ways*shots
    adapt_idx = np.zeros(X.size(0), dtype=bool)
    adapt_idx[np.arange(1*5) * 2] = True # shots*ways
    eval_idx = torch.from_numpy(~adapt_idx)
    apt_idx = torch.from_numpy(adapt_idx)
    adapt_data, adapt_labels = X[adapt_idx], y[adapt_idx]
    eval_data, eval_labels = X[eval_idx], y[eval_idx]

    # Adaptation Steps (Steps: 0 to K-1)
    for steps in range(0, steps): # steps : no of adaption steps
        pred = learner(adapt_data)
        loss = loss_fn(pred, adapt_labels)
        grad = torch.autograd.grad(loss, learner.parameters())
        new_weight = list(map(lambda p: p[1] - 0.4*p[0], zip(grad, learner.parameters()))) # theta' = theta - alpha*grads
        with torch.no_grad(): # Updating the new_weights to learner
            for i, params in enumerate(learner.parameters()):
                params.copy_(new_weight[i])
        
    # Evaluation loss, accuracy
    pred_val = learner(eval_data)
    loss_val = loss_fn(pred_val, eval_labels)
    pred_val = pred_val.argmax(dim=1).view(eval_labels.shape)
    acc_val = (pred_val == eval_labels).sum().float()/eval_labels.size(0)
    
    return loss_val, acc_val
</code></pre>
<p>And the outer loop:</p>
<pre><code class="lang-auto">iterations = 5
meta_batch_size = 32
for iteration in range(iterations):
    train_loss, train_acc = 0, 0
    for task in range(meta_batch_size):
        # Copy model parameters. 
        learner = deepcopy(new_model)
        batch = tasksets.train.sample()  # Generates 2*shots*ways samples.
        # The batch will be split into Dtrain and Dtest in the inner loop (both of size shots*ways)
        loss_val, acc_val = inner_loop(learner, batch, loss_fn, 1, device) # Gradient steps: 1
        train_loss += loss_val/meta_batch_size
        train_acc += acc_val/meta_batch_size
        del learner
    meta_opt.zero_grad()
    train_loss.backward()
    meta_opt.step()
    if iteration % 1 == 0:
        print('Training Accuracy:', train_acc.item())
        print('Training Loss:', train_loss.item())
</code></pre>
<p>As you might have guessed, the computation graph is breaking in the outer loop, thus <code>train_loss</code> is not updating any weights. It would be awesome if you are able to point out a few things! Thank you!</p>
<p>Just for the sake of additional information, I will add the Train Loss, Acc here:</p>
<pre><code class="lang-auto">Training Accuracy: 0.3437499403953552
Training Loss: 1.4747799634933472
Training Accuracy: 0.3062499761581421
Training Loss: 1.510756254196167
Training Accuracy: 0.38749998807907104
Training Loss: 1.5028401613235474
Training Accuracy: 0.3312499225139618
Training Loss: 1.4998550415039062
Training Accuracy: 0.35625001788139343
Training Loss: 1.477844476699829
</code></pre>

Tushar_Prakash: <p>Hello <a class="mention" href="/u/asura">@asura</a>, Have you resolved this issue?</p>

asura: <p>Hey <a class="mention" href="/u/tushar_prakash">@Tushar_Prakash</a> what problem are you facing? It’s been quite a while since I worked on this, but I will try to help as much as I can.</p>

Tushar_Prakash: <p>Hello <a class="mention" href="/u/asura">@asura</a> thanks for your reply. I just saw the last reply date.</p>
<p>It’s same as you described above. So, when we copy the original model weights for the inner loop, update them, and then compute the loss by running a forward pass on the query set with the updated model, we hit a snag. Specifically, when we attempt to obtain the gradients with respect to the original weights, it throws an error because the original weights aren’t included in the computation graph. How did you manage to resolve this issue?</p>