Thanks a lot for pointing me this library. It made things a lot easier.
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.
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')

Just in case, full code:

  
      github.com
  
  
    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 10
{
 "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",


  This file has been truncated. show original