Autograd doesn't retain the computation graph on a meta-learning algorithm

TinfoilHat0: <p>Hi,</p>
<p>I’m trying to implement the algorithm described <a href="https://arxiv.org/abs/1803.09050" rel="noopener nofollow ugc">here.</a><br>
A short description of the related part is as follows.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/5/b57c6c089296ba2147aa73d9c884852e440ef74c.png" data-download-href="https://discuss.pytorch.org/uploads/default/b57c6c089296ba2147aa73d9c884852e440ef74c" title="Screenshot from 2020-04-22 06-09-37"><img src="https://discuss.pytorch.org/uploads/default/original/3X/b/5/b57c6c089296ba2147aa73d9c884852e440ef74c.png" alt="Screenshot from 2020-04-22 06-09-37" data-base62-sha1="pTuZH2SkeHnZtdXB2pa8HTGnzZO" width="503" height="500" data-dominant-color="F0F0F0"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot from 2020-04-22 06-09-37</span><span class="informations">527×523 48.5 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>My attempt is as follows.</p>
<pre><code class="lang-auto"># 1. Forward-backward pass on training data
_, (inputs, labels) =  next(enumerate(train_loader))
inputs, labels = inputs.to(device=args.device, non_blocking=True),\
                        labels.to(device=args.device, non_blocking=True)
meta_model.load_state_dict(model.state_dict())

y_hat_f = meta_model(inputs)
criterion.reduction = 'none'
l_f = criterion(y_hat_f, labels)
eps = torch.rand(l_f.size(), requires_grad=False, device=args.device).div(1e6)
eps.requires_grad = True
l_f = torch.sum(eps * l_f)

# 2. Compute grads wrt model and update its params
l_f.backward(retain_graph=True)
meta_optimizer.step()

# 3. Forward-backward pass on meta data with updated model
_, (inputs, labels) =  next(enumerate(meta_loader))
inputs, labels = inputs.to(device=args.device, non_blocking=True),\
                        labels.to(device=args.device, non_blocking=True)

y_hat_g = model(inputs)
criterion.reduction = 'mean'
l_g = criterion(y_hat_g, labels)

# 4. Compute grads wrt eps and update weights
eps_grads = torch.autograd.grad(l_g, eps)
.....
</code></pre>
<p>At this line:</p>
<blockquote>
<p>eps_grads = torch.autograd.grad(l_g, eps)</p>
</blockquote>
<p>I get an error saying,</p>
<blockquote>
<p>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.</p>
</blockquote>
<p>If I set <strong>allow_used=True</strong>, it returns <em><strong>None</strong></em> to <strong>eps_grad.</strong></p>
<p>As far as I can tell, autograd loses computation graph for some reason and doesn’t retain the information that <strong>eps</strong> was used in the computation of <strong>l_f</strong>, which in turn used in updating parameters of the model.  So, <strong>l_g</strong> should be differentiable wrt <strong>eps</strong> but it doesn’t work here.</p>
<p>How can I solve this ?</p>

albanD: <p>Hi,</p>
<p>I think you’re missing a <code>create_graph=True</code> in the first backward if you want to be able to backprop through the backward pass.<br>
Also pytorch’s optimizers are not differentiable (yet). So you won’t be able to backprop through that step either.</p>
<p>I would recommend using the higher package to do this: <a href="https://github.com/facebookresearch/higher">https://github.com/facebookresearch/higher</a></p>

TinfoilHat0: <p>so i should update the parameters manually, presumably looping through them and deducting gradients, given optimizers are not differentiable?</p>

albanD: <p>Yes,<br>
But the tricky bit is that <code>nn.Parameter()</code> are built to be parameters that you learn. So they cannot have history. So you will have to delete these and replace them with the new updated values as Tensors (and keep them in a different place so that you can still update them with your optimizer).<br>
That is why I recommended the library above that does all that for you <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

TinfoilHat0: <p>Gotcha. Thank a lot.</p>

TinfoilHat0: <p>Thanks a lot for pointing me this library. It made things a lot easier.</p>
<p>I’ve a question regarding the performance though. Roughly, the code is 4-5 times slower than what it used to be. Given meta learning algorithm does 3x forward-backward passes, it seems sort of OK to me. Just in case, could you take a quick look and tell me whether I’m doing something grossly inefficient in the following code ? Also, I’ve been able to replicate an experiment from the paper so I believe it is OK from the correctness perspective.</p>
<pre><code class="lang-auto">for rnd in tqdm(range(1, args.epochs+1)):
    model.train()
    train_loss, train_acc = 0.0, 0.0 
    
    for _, (inputs, labels) in enumerate(train_loader):
        inputs, labels = inputs.to(device=args.device, non_blocking=True),\
                        labels.to(device=args.device, non_blocking=True)
        opt.zero_grad()
        
        with higher.innerloop_ctx(model, opt) as (meta_model, meta_opt):
            # 1. Update meta model on training data
            meta_train_outputs = meta_model(inputs)
            criterion.reduction = 'none'
            meta_train_loss = criterion(meta_train_outputs, labels)
            eps = torch.rand(meta_train_loss.size(), requires_grad=False, device=args.device).div(1e6)
            eps.requires_grad = True
            meta_train_loss = torch.sum(eps * meta_train_loss)
            meta_opt.step(meta_train_loss)
            
            # 2. Compute grads of eps on meta validation data
            meta_val_outputs = meta_model(meta_inputs)
            criterion.reduction = 'mean'
            meta_val_loss = criterion(meta_val_outputs, meta_labels)
            eps_grads = torch.autograd.grad(meta_val_loss, eps, only_inputs=True)[0].detach()
            
        # 3. Compute weights for current training batch
        w_tilde = torch.clamp(-eps_grads, min=0)
        l1_norm = torch.sum(w_tilde)
        if l1_norm != 0:
            w = w_tilde / l1_norm
        else:
            w = w_tilde
            
        # 4. Train model on weighted batch
        outputs = model(inputs)
        criterion.reduction = 'none'
        minibatch_loss = criterion(outputs, labels)
        minibatch_loss = torch.sum(w * minibatch_loss)
        minibatch_loss.backward()
        
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=5)
        opt.step()
        
        # keep track of epoch loss/accuracy
        train_loss += minibatch_loss.item()*outputs.shape[0]
        _, pred_labels = torch.max(outputs, 1)
        train_acc += torch.sum(torch.eq(pred_labels.view(-1), labels)).item()
        
    # inference after epoch
    with torch.no_grad():
        train_loss, train_acc = train_loss/len(train_dataset), train_acc/len(train_dataset)       
        val_loss, (val_acc, val_per_class) = utils.get_loss_n_accuracy(model, criterion, val_loader, args)                                  
        scheduler.step(val_loss)
        # log/print data
        #writer.add_scalar('Validation/Loss', val_loss, rnd)
        #writer.add_scalar('Validation/Accuracy', val_acc, rnd)
        #writer.add_scalar('Training/Loss', train_loss, rnd)
        #writer.add_scalar('Training/Accuracy', train_acc, rnd)
        print(f'|Train/Valid Loss: {train_loss:.3f} / {val_loss:.3f}|', end='--')
        print(f'|Train/Valid Acc: {train_acc:.3f} / {val_acc:.3f}|', end='\r')
</code></pre>
<p>Just in case, full code:</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/TinfoilHat0/Learning-to-Reweight-Examples-for-Robust-Deep-Learning-with-PyTorch-Higher/blob/master/Learning%20to%20Reweight%20Examples%20for%20Robust%20Deep%20Learning%20with%20PyTorch%20Higher.ipynb" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/TinfoilHat0/Learning-to-Reweight-Examples-for-Robust-Deep-Learning-with-PyTorch-Higher/blob/master/Learning%20to%20Reweight%20Examples%20for%20Robust%20Deep%20Learning%20with%20PyTorch%20Higher.ipynb" target="_blank" rel="nofollow noopener">TinfoilHat0/Learning-to-Reweight-Examples-for-Robust-Deep-Learning-with-PyTorch-Higher/blob/master/Learning to Reweight Examples for Robust Deep Learning with PyTorch Higher.ipynb</a></h4>
<pre><code class="lang-ipynb">{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.insert(1, '../src/')\n",
    "import torch\n",
    "from tqdm.notebook import tqdm\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.optim import lr_scheduler\n",
    "import numpy as np\n",
    "import torchvision\n",
    "from torchvision import datasets,transforms\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "from torch.utils.data import DataLoader\n",
</code></pre>

  This file has been truncated. <a href="https://github.com/TinfoilHat0/Learning-to-Reweight-Examples-for-Robust-Deep-Learning-with-PyTorch-Higher/blob/master/Learning%20to%20Reweight%20Examples%20for%20Robust%20Deep%20Learning%20with%20PyTorch%20Higher.ipynb" target="_blank" rel="nofollow noopener">show original</a>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>


albanD: <p>Your old code was not running any backward. So it was expected that this will be slower <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=9" title=":smiley:" class="emoji" alt=":smiley:"></p>
<p>The code looks ok to me.</p>