cifar/cifar10_defer_baselines.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook implements the described baselines for our CIFAR experiments: LearnedOracle, Confidence and Oracle.\n",
    "\n",
    "Running this notebook will produce results for a single expert k for the 3 baselines."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "g1mubTgqeM88"
   },
   "source": [
    "# Load data and model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "0V9WxRqMd-Al",
    "outputId": "b861708b-f547-40b7-dbf0-fa787a00350c"
   },
   "outputs": [],
   "source": [
    "import math\n",
    "import torch\n",
    "import numpy as np\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import argparse\n",
    "import os\n",
    "import random\n",
    "import shutil\n",
    "import time\n",
    "import torch.nn.parallel\n",
    "import torch.backends.cudnn as cudnn\n",
    "import torch.optim\n",
    "import torch.utils.data\n",
    "import torchvision.transforms as transforms\n",
    "import torchvision.datasets as datasets\n",
    "from torch.autograd import Variable\n",
    "\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(device)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "D8b5SxX0eDQQ"
   },
   "outputs": [],
   "source": [
    "class BasicBlock(nn.Module):\n",
    "    def __init__(self, in_planes, out_planes, stride, dropRate=0.0):\n",
    "        super(BasicBlock, self).__init__()\n",
    "        self.bn1 = nn.BatchNorm2d(in_planes)\n",
    "        self.relu1 = nn.ReLU(inplace=True)\n",
    "        self.conv1 = nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,\n",
    "                               padding=1, bias=False)\n",
    "        self.bn2 = nn.BatchNorm2d(out_planes)\n",
    "        self.relu2 = nn.ReLU(inplace=True)\n",
    "        self.conv2 = nn.Conv2d(out_planes, out_planes, kernel_size=3, stride=1,\n",
    "                               padding=1, bias=False)\n",
    "        self.droprate = dropRate\n",
    "        self.equalInOut = (in_planes == out_planes)\n",
    "        self.convShortcut = (not self.equalInOut) and nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride,\n",
    "                                                                padding=0, bias=False) or None\n",
    "    def forward(self, x):\n",
    "        if not self.equalInOut:\n",
    "            x = self.relu1(self.bn1(x))\n",
    "        else:\n",
    "            out = self.relu1(self.bn1(x))\n",
    "        out = self.relu2(self.bn2(self.conv1(out if self.equalInOut else x)))\n",
    "        if self.droprate > 0:\n",
    "            out = F.dropout(out, p=self.droprate, training=self.training)\n",
    "        out = self.conv2(out)\n",
    "        return torch.add(x if self.equalInOut else self.convShortcut(x), out)\n",
    "\n",
    "\n",
    "class NetworkBlock(nn.Module):\n",
    "    def __init__(self, nb_layers, in_planes, out_planes, block, stride, dropRate=0.0):\n",
    "        super(NetworkBlock, self).__init__()\n",
    "        self.layer = self._make_layer(block, in_planes, out_planes, nb_layers, stride, dropRate)\n",
    "    def _make_layer(self, block, in_planes, out_planes, nb_layers, stride, dropRate):\n",
    "        layers = []\n",
    "        for i in range(int(nb_layers)):\n",
    "            layers.append(block(i == 0 and in_planes or out_planes, out_planes, i == 0 and stride or 1, dropRate))\n",
    "        return nn.Sequential(*layers)\n",
    "    def forward(self, x):\n",
    "        return self.layer(x)\n",
    "\n",
    "\n",
    "class WideResNet(nn.Module):\n",
    "    def __init__(self, depth, num_classes, widen_factor=1, dropRate=0.0):\n",
    "        super(WideResNet, self).__init__()\n",
    "        nChannels = [16, 16 * widen_factor, 32 * widen_factor, 64 * widen_factor]\n",
    "        assert ((depth - 4) % 6 == 0)\n",
    "        n = (depth - 4) / 6\n",
    "        block = BasicBlock\n",
    "        # 1st conv before any network block\n",
    "        self.conv1 = nn.Conv2d(3, nChannels[0], kernel_size=3, stride=1,\n",
    "                               padding=1, bias=False)\n",
    "        # 1st block\n",
    "        self.block1 = NetworkBlock(n, nChannels[0], nChannels[1], block, 1, dropRate)\n",
    "        # 2nd block\n",
    "        self.block2 = NetworkBlock(n, nChannels[1], nChannels[2], block, 2, dropRate)\n",
    "        # 3rd block\n",
    "        self.block3 = NetworkBlock(n, nChannels[2], nChannels[3], block, 2, dropRate)\n",
    "        # global average pooling and classifier\n",
    "        self.bn1 = nn.BatchNorm2d(nChannels[3])\n",
    "        self.relu = nn.ReLU(inplace=True)\n",
    "        self.fc = nn.Linear(nChannels[3], num_classes)\n",
    "        self.nChannels = nChannels[3]\n",
    "        self.softmax = nn.Softmax()\n",
    "        for m in self.modules():\n",
    "            if isinstance(m, nn.Conv2d):\n",
    "                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')\n",
    "            elif isinstance(m, nn.BatchNorm2d):\n",
    "                m.weight.data.fill_(1)\n",
    "                m.bias.data.zero_()\n",
    "            elif isinstance(m, nn.Linear):\n",
    "                m.bias.data.zero_()\n",
    "    def forward(self, x):\n",
    "        out = self.conv1(x)\n",
    "        out = self.block1(out)\n",
    "        out = self.block2(out)\n",
    "        out = self.block3(out)\n",
    "        out = self.relu(self.bn1(out))\n",
    "        out = F.avg_pool2d(out, 8)\n",
    "        out = out.view(-1, self.nChannels)\n",
    "        out = self.fc(out)\n",
    "        out = self.softmax(out)\n",
    "        return out\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "NzGKUz5LfHHI"
   },
   "outputs": [],
   "source": [
    "class AverageMeter(object):\n",
    "    \"\"\"Computes and stores the average and current value\"\"\"\n",
    "    def __init__(self):\n",
    "        self.reset()\n",
    "\n",
    "    def reset(self):\n",
    "        self.val = 0\n",
    "        self.avg = 0\n",
    "        self.sum = 0\n",
    "        self.count = 0\n",
    "\n",
    "    def update(self, val, n=1):\n",
    "        self.val = val\n",
    "        self.sum += val * n\n",
    "        self.count += n\n",
    "        self.avg = self.sum / self.count\n",
    "\n",
    "def accuracy(output, target, topk=(1,)):\n",
    "    \"\"\"Computes the precision@k for the specified values of k\"\"\"\n",
    "    maxk = max(topk)\n",
    "    batch_size = target.size(0)\n",
    "\n",
    "    _, pred = output.topk(maxk, 1, True, True)\n",
    "    pred = pred.t()\n",
    "    correct = pred.eq(target.view(1, -1).expand_as(pred))\n",
    "\n",
    "    res = []\n",
    "    for k in topk:\n",
    "        correct_k = correct[:k].view(-1).float().sum(0)\n",
    "        res.append(correct_k.mul_(100.0 / batch_size))\n",
    "    return res\n",
    "\n",
    "\n",
    "def metrics_print(net, expert_fn, n_classes, loader):\n",
    "    '''\n",
    "    Computes metrics for deferal\n",
    "    -----\n",
    "    Arguments:\n",
    "    net: model\n",
    "    expert_fn: expert model predict function\n",
    "    n_classes: number of classes\n",
    "    loader: data loader\n",
    "    '''\n",
    "    correct = 0\n",
    "    correct_sys = 0\n",
    "    exp = 0\n",
    "    exp_total = 0\n",
    "    total = 0\n",
    "    real_total = 0\n",
    "    with torch.no_grad():\n",
    "        for data in loader:\n",
    "            images, labels = data\n",
    "            images, labels = images.to(device), labels.to(device)\n",
    "            outputs = net(images)\n",
    "            _, predicted = torch.max(outputs.data, 1)\n",
    "            batch_size = outputs.size()[0]  # batch_size\n",
    "            exp_prediction = expert_fn(images, labels)\n",
    "            for i in range(0, batch_size):\n",
    "                r = (predicted[i].item() == n_classes)\n",
    "                if r == 0:\n",
    "                    total += 1\n",
    "                    correct += (predicted[i] == labels[i]).item()\n",
    "                    correct_sys += (predicted[i] == labels[i]).item()\n",
    "                if r == 1:\n",
    "                    exp += (exp_prediction[i] == labels[i].item())\n",
    "                    correct_sys += (exp_prediction[i] == labels[i].item())\n",
    "                    exp_total += 1\n",
    "                real_total += 1\n",
    "    cov = str(total) + str(\" out of\") + str(real_total)\n",
    "    to_print = {\"coverage\": cov, \"system accuracy\": 100 * correct_sys / real_total,\n",
    "                \"expert accuracy\": 100 * exp / (exp_total + 0.0002),\n",
    "                \"classifier accuracy\": 100 * correct / (total + 0.0001)}\n",
    "    print(to_print)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "7iYuKa2KKap9"
   },
   "source": [
    "# Baseline: LearnedOracle\n",
    "Classifier: trained on all the data using cross entropy\n",
    "\n",
    "Rejector: trained to recognize if image is in the first k classes (with knowledge of k)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "XgenUJHrLLRn"
   },
   "source": [
    "## Learn classifier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "_2DwWykXKdq2"
   },
   "outputs": [],
   "source": [
    "def my_CrossEntropyLoss(outputs, labels):\n",
    "    batch_size = outputs.size()[0]  # batch_size\n",
    "    outputs = - torch.log2(outputs[range(batch_size), labels])  # pick the values corresponding to the labels\n",
    "    return torch.sum(outputs) / batch_size\n",
    "\n",
    "def train_classifier(train_loader, model, optimizer, scheduler, epoch, expert_fn, n_classes):\n",
    "    \"\"\"Train for one epoch on the training set\"\"\"\n",
    "    # expertfn: a number here k \n",
    "    batch_time = AverageMeter()\n",
    "    losses = AverageMeter()\n",
    "    top1 = AverageMeter()\n",
    "\n",
    "    # switch to train mode\n",
    "    model.train()\n",
    "\n",
    "    end = time.time()\n",
    "    for i, (input, target) in enumerate(train_loader):\n",
    "        target = target.to(device)\n",
    "        input = input.to(device)\n",
    "\n",
    "        # compute output\n",
    "        output = model(input)\n",
    "        # compute loss\n",
    "        loss = my_CrossEntropyLoss(output, target)\n",
    "\n",
    "        # measure accuracy and record loss\n",
    "        prec1 = accuracy(output.data, target, topk=(1,))[0]\n",
    "        losses.update(loss.data.item(), input.size(0))\n",
    "        top1.update(prec1.item(), input.size(0))\n",
    "\n",
    "        # compute gradient and do SGD step\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        scheduler.step()\n",
    "\n",
    "        # measure elapsed time\n",
    "        batch_time.update(time.time() - end)\n",
    "        end = time.time()\n",
    "\n",
    "        if i % 10 == 0:\n",
    "            print('Epoch: [{0}][{1}/{2}]\\t'\n",
    "                  'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\\t'\n",
    "                  'Loss {loss.val:.4f} ({loss.avg:.4f})\\t'\n",
    "                  'Prec@1 {top1.val:.3f} ({top1.avg:.3f})'.format(\n",
    "                epoch, i, len(train_loader), batch_time=batch_time,\n",
    "                loss=losses, top1=top1))\n",
    "\n",
    "\n",
    "def validate_classifier(val_loader, model, epoch, expert_fn, n_classes):\n",
    "    \"\"\"Perform validation on the validation set\"\"\"\n",
    "    batch_time = AverageMeter()\n",
    "    losses = AverageMeter()\n",
    "    top1 = AverageMeter()\n",
    "\n",
    "    # switch to evaluate mode\n",
    "    model.eval()\n",
    "\n",
    "    end = time.time()\n",
    "    for i, (input, target) in enumerate(val_loader):\n",
    "        target = target.to(device)\n",
    "        input = input.to(device)\n",
    "\n",
    "        # compute output\n",
    "        with torch.no_grad():\n",
    "            output = model(input)\n",
    "        # compute loss\n",
    "        loss = my_CrossEntropyLoss(output, target)\n",
    "\n",
    "        # measure accuracy and record loss\n",
    "        prec1 = accuracy(output.data, target, topk=(1,))[0]\n",
    "        losses.update(loss.data.item(), input.size(0))\n",
    "        top1.update(prec1.item(), input.size(0))\n",
    "\n",
    "        # measure elapsed time\n",
    "        batch_time.update(time.time() - end)\n",
    "        end = time.time()\n",
    "\n",
    "        if i % 10 == 0:\n",
    "            print('Test: [{0}/{1}]\\t'\n",
    "                  'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\\t'\n",
    "                  'Loss {loss.val:.4f} ({loss.avg:.4f})\\t'\n",
    "                  'Prec@1 {top1.val:.3f} ({top1.avg:.3f})'.format(\n",
    "                i, len(val_loader), batch_time=batch_time, loss=losses,\n",
    "                top1=top1))\n",
    "\n",
    "    print(' * Prec@1 {top1.avg:.3f}'.format(top1=top1))\n",
    "\n",
    "    return top1.avg\n",
    "best_prec1 = 0\n",
    "def run_classifier(model, data_aug, n_dataset, expert_fn, epochs):\n",
    "    global best_prec1\n",
    "    # Data loading code\n",
    "    normalize = transforms.Normalize(mean=[x / 255.0 for x in [125.3, 123.0, 113.9]],\n",
    "                                     std=[x / 255.0 for x in [63.0, 62.1, 66.7]])\n",
    "\n",
    "    if data_aug:\n",
    "        transform_train = transforms.Compose([\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Lambda(lambda x: F.pad(x.unsqueeze(0),\n",
    "                                              (4, 4, 4, 4), mode='reflect').squeeze()),\n",
    "            transforms.ToPILImage(),\n",
    "            transforms.RandomCrop(32),\n",
    "            transforms.RandomHorizontalFlip(),\n",
    "            transforms.ToTensor(),\n",
    "            normalize,\n",
    "        ])\n",
    "    else:\n",
    "        transform_train = transforms.Compose([\n",
    "            transforms.ToTensor(),\n",
    "            normalize,\n",
    "        ])\n",
    "    transform_test = transforms.Compose([\n",
    "        transforms.ToTensor(),\n",
    "        normalize\n",
    "    ])\n",
    "\n",
    "    if n_dataset == 10:\n",
    "        dataset = 'cifar10'\n",
    "    elif n_dataset == 100:\n",
    "        dataset = 'cifar100'\n",
    "\n",
    "    kwargs = {'num_workers': 0, 'pin_memory': True}\n",
    "\n",
    "    train_dataset_all = datasets.__dict__[dataset.upper()]('../data', train=True, download=True,\n",
    "                                                           transform=transform_train)\n",
    "    train_size = int(0.90 * len(train_dataset_all))\n",
    "    test_size = len(train_dataset_all) - train_size\n",
    "\n",
    "    train_dataset, test_dataset = torch.utils.data.random_split(train_dataset_all, [train_size, test_size])\n",
    "    train_loader = torch.utils.data.DataLoader(train_dataset,\n",
    "                                               batch_size=128, shuffle=True, **kwargs)\n",
    "    test_loader = torch.utils.data.DataLoader(test_dataset,\n",
    "                                               batch_size=128, shuffle=True, **kwargs)\n",
    "\n",
    "    # get the number of model parameters\n",
    "    print('Number of model parameters: {}'.format(\n",
    "        sum([p.data.nelement() for p in model.parameters()])))\n",
    "\n",
    "    # for training on multiple GPUs.\n",
    "    # Use CUDA_VISIBLE_DEVICES=0,1 to specify which GPUs to use\n",
    "    # model = torch.nn.DataParallel(model).cuda()\n",
    "    model = model.to(device)\n",
    "\n",
    "    # optionally resume from a checkpoint\n",
    "\n",
    "    cudnn.benchmark = True\n",
    "\n",
    "    # define loss function (criterion) and optimizer\n",
    "    optimizer = torch.optim.SGD(model.parameters(), 0.1,\n",
    "                                momentum=0.9, nesterov=True,\n",
    "                                weight_decay=5e-4)\n",
    "\n",
    "    # cosine learning rate\n",
    "    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, len(train_loader) * 200)\n",
    "\n",
    "    for epoch in range(0, epochs):\n",
    "        # train for one epoch\n",
    "        train_classifier(train_loader, model, optimizer, scheduler, epoch, expert_fn, n_dataset)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "M1LZykRYK1WE"
   },
   "outputs": [],
   "source": [
    "n_dataset = 10  # cifar-10, 100 for cifar-100\n",
    "model_classifier = WideResNet(28, n_dataset, 4, dropRate=0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "6jx5f_DBK-ek"
   },
   "outputs": [],
   "source": [
    "run_classifier(model_classifier, False, n_dataset, 0, 200)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "OeR3zUBZLTVK"
   },
   "source": [
    "## Train rejector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "qQBq5Su7LUh3"
   },
   "outputs": [],
   "source": [
    "def my_CrossEntropyLoss(outputs, labels):\n",
    "    # m: expert costs, labels: ground truth, n_classes: number of classes\n",
    "    batch_size = outputs.size()[0]  # batch_size\n",
    "    outputs = - torch.log2(outputs[range(batch_size), labels])  # pick the values corresponding to the labels\n",
    "    return torch.sum(outputs) / batch_size\n",
    "\n",
    "def train_rejector(train_loader, model, optimizer, scheduler, epoch, expert_fn, n_classes):\n",
    "    \"\"\"Train for one epoch on the training set\"\"\"\n",
    "    batch_time = AverageMeter()\n",
    "    losses = AverageMeter()\n",
    "    top1 = AverageMeter()\n",
    "    #expertfn is k (a number)\n",
    "    # switch to train mode\n",
    "    model.train()\n",
    "\n",
    "    end = time.time()\n",
    "    for i, (input, target) in enumerate(train_loader):\n",
    "        target = target.to(device)\n",
    "        input = input.to(device)\n",
    "\n",
    "        # compute output\n",
    "        output = model(input)\n",
    "        # compute new target\n",
    "        batch_size = output.size()[0]  # batch_size\n",
    "        m = [0] * batch_size\n",
    "        for j in range(0, batch_size):\n",
    "            if target[j].item() <= expert_fn:\n",
    "                m[j] = 1\n",
    "            else:\n",
    "                m[j] = 0\n",
    "        m = torch.tensor(m)\n",
    "        m = m.to(device)\n",
    "        # compute loss\n",
    "        loss = my_CrossEntropyLoss(output, m)\n",
    "\n",
    "        # measure accuracy and record loss\n",
    "        prec1 = accuracy(output.data, m, topk=(1,))[0]\n",
    "        losses.update(loss.data.item(), input.size(0))\n",
    "        top1.update(prec1.item(), input.size(0))\n",
    "\n",
    "        # compute gradient and do SGD step\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        scheduler.step()\n",
    "\n",
    "        # measure elapsed time\n",
    "        batch_time.update(time.time() - end)\n",
    "        end = time.time()\n",
    "\n",
    "        if i % 10 == 0:\n",
    "            print('Epoch: [{0}][{1}/{2}]\\t'\n",
    "                  'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\\t'\n",
    "                  'Loss {loss.val:.4f} ({loss.avg:.4f})\\t'\n",
    "                  'Prec@1 {top1.val:.3f} ({top1.avg:.3f})'.format(\n",
    "                epoch, i, len(train_loader), batch_time=batch_time,\n",
    "                loss=losses, top1=top1))\n",
    "\n",
    "\n",
    "def validate_rejector(val_loader, model, epoch, expert_fn, n_classes):\n",
    "    \"\"\"Perform validation on the validation set\"\"\"\n",
    "    batch_time = AverageMeter()\n",
    "    losses = AverageMeter()\n",
    "    top1 = AverageMeter()\n",
    "\n",
    "    # switch to evaluate mode\n",
    "    model.eval()\n",
    "\n",
    "    end = time.time()\n",
    "    for i, (input, target) in enumerate(val_loader):\n",
    "        target = target.to(device)\n",
    "        input = input.to(device)\n",
    "\n",
    "        # compute output\n",
    "        with torch.no_grad():\n",
    "            output = model(input)\n",
    "        # expert prediction\n",
    "        batch_size = output.size()[0]  # batch_size\n",
    "        m = [0] * batch_size\n",
    "        for j in range(0, batch_size):\n",
    "            if target[j].item() <= expert_fn:\n",
    "                m[j] = 1\n",
    "            else:\n",
    "                m[j] = 0\n",
    "        m = torch.tensor(m)\n",
    "        m = m.to(device)\n",
    "        # compute loss\n",
    "        loss = my_CrossEntropyLoss(output, m)\n",
    "\n",
    "        # measure accuracy and record loss\n",
    "        prec1 = accuracy(output.data, m, topk=(1,))[0]\n",
    "        losses.update(loss.data.item(), input.size(0))\n",
    "        top1.update(prec1.item(), input.size(0))\n",
    "\n",
    "        # measure elapsed time\n",
    "        batch_time.update(time.time() - end)\n",
    "        end = time.time()\n",
    "\n",
    "        if i % 10 == 0:\n",
    "            print('Test: [{0}/{1}]\\t'\n",
    "                  'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\\t'\n",
    "                  'Loss {loss.val:.4f} ({loss.avg:.4f})\\t'\n",
    "                  'Prec@1 {top1.val:.3f} ({top1.avg:.3f})'.format(\n",
    "                i, len(val_loader), batch_time=batch_time, loss=losses,\n",
    "                top1=top1))\n",
    "\n",
    "    print(' * Prec@1 {top1.avg:.3f}'.format(top1=top1))\n",
    "\n",
    "    return top1.avg\n",
    "\n",
    "def run_rejector(model, data_aug, n_dataset, expert_fn, epochs):\n",
    "    global best_prec1\n",
    "    # Data loading code\n",
    "    normalize = transforms.Normalize(mean=[x / 255.0 for x in [125.3, 123.0, 113.9]],\n",
    "                                     std=[x / 255.0 for x in [63.0, 62.1, 66.7]])\n",
    "\n",
    "    if data_aug:\n",
    "        transform_train = transforms.Compose([\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Lambda(lambda x: F.pad(x.unsqueeze(0),\n",
    "                                              (4, 4, 4, 4), mode='reflect').squeeze()),\n",
    "            transforms.ToPILImage(),\n",
    "            transforms.RandomCrop(32),\n",
    "            transforms.RandomHorizontalFlip(),\n",
    "            transforms.ToTensor(),\n",
    "            normalize,\n",
    "        ])\n",
    "    else:\n",
    "        transform_train = transforms.Compose([\n",
    "            transforms.ToTensor(),\n",
    "            normalize,\n",
    "        ])\n",
    "    transform_test = transforms.Compose([\n",
    "        transforms.ToTensor(),\n",
    "        normalize\n",
    "    ])\n",
    "\n",
    "    if n_dataset == 10:\n",
    "        dataset = 'cifar10'\n",
    "    elif n_dataset == 100:\n",
    "        dataset = 'cifar100'\n",
    "\n",
    "    kwargs = {'num_workers': 0, 'pin_memory': True}\n",
    "\n",
    "    train_dataset_all = datasets.__dict__[dataset.upper()]('../data', train=True, download=True,\n",
    "                                                           transform=transform_train)\n",
    "    train_size = int(0.90 * len(train_dataset_all))\n",
    "    test_size = len(train_dataset_all) - train_size\n",
    "\n",
    "    train_dataset, test_dataset = torch.utils.data.random_split(train_dataset_all, [train_size, test_size])\n",
    "    train_loader = torch.utils.data.DataLoader(train_dataset,\n",
    "                                               batch_size=128, shuffle=True, **kwargs)\n",
    "    test_loader = torch.utils.data.DataLoader(test_dataset,\n",
    "                                               batch_size=128, shuffle=True, **kwargs)\n",
    "\n",
    "    # get the number of model parameters\n",
    "    print('Number of model parameters: {}'.format(\n",
    "        sum([p.data.nelement() for p in model.parameters()])))\n",
    "\n",
    "    # for training on multiple GPUs.\n",
    "    # Use CUDA_VISIBLE_DEVICES=0,1 to specify which GPUs to use\n",
    "    # model = torch.nn.DataParallel(model).cuda()\n",
    "    model = model.to(device)\n",
    "\n",
    "    # optionally resume from a checkpoint\n",
    "\n",
    "    cudnn.benchmark = True\n",
    "\n",
    "    # define loss function (criterion) and optimizer\n",
    "    optimizer = torch.optim.SGD(model.parameters(), 0.1,\n",
    "                                momentum=0.9, nesterov=True,\n",
    "                                weight_decay=5e-4)\n",
    "\n",
    "    # cosine learning rate\n",
    "    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, len(train_loader) * 200)\n",
    "\n",
    "    for epoch in range(0, epochs):\n",
    "        # train for one epoch\n",
    "        train_rejector(train_loader, model, optimizer, scheduler, epoch, expert_fn, n_dataset)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "k below is the number of classes the exper can predict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "k = 5  # max k is 9 and min k is 0\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ijL7CU2YLtyg"
   },
   "outputs": [],
   "source": [
    "class synth_expert:\n",
    "    def __init__(self, k, n_classes):\n",
    "        self.k = k\n",
    "        self.n_classes = n_classes\n",
    "\n",
    "    def predict(self, input, labels):\n",
    "        batch_size = labels.size()[0]  # batch_size\n",
    "        outs = [0] * batch_size\n",
    "        for i in range(0, batch_size):\n",
    "            if labels[i].item() <= self.k:\n",
    "                outs[i] = labels[i].item()\n",
    "            else:\n",
    "                prediction_rand = random.randint(0, self.n_classes - 1)\n",
    "                outs[i] = prediction_rand\n",
    "        return outs\n",
    "\n",
    "\n",
    "expert = synth_expert(k, n_dataset)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "YW_Vchg7zbeH"
   },
   "outputs": [],
   "source": [
    "model_rejector = WideResNet(10, 2, 4, dropRate=0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "colab_type": "code",
    "id": "gzpQDai4X6R_",
    "outputId": "27e43f87-d4ae-4bda-c7f5-7fc7e7c88f36"
   },
   "outputs": [],
   "source": [
    "run_rejector(model_rejector, False, n_dataset, k, 200)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "EmJzEurLML0y"
   },
   "source": [
    "## Evaluate\n",
    "This produces the key metrics measured:\n",
    "- coverage: percentage of examples where classifier predicts\n",
    "- system accuracy: accuracy of combined system on all data (with deferral)\n",
    "- classifier accuracy: accuracy of classifier on non-deferred examples\n",
    "- expert accuracy: accuracy of expert on deferred examples\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 92
    },
    "colab_type": "code",
    "id": "pCIIsQA2MNbW",
    "outputId": "5034b52a-49fb-4e12-afe5-9b5dac37f2dd"
   },
   "outputs": [],
   "source": [
    "def metrics_print_baseline(net_class, net_rej, expert_fn, n_classes, loader):\n",
    "    correct = 0\n",
    "    correct_sys = 0\n",
    "    exp = 0\n",
    "    exp_total = 0\n",
    "    total = 0\n",
    "    real_total = 0\n",
    "    with torch.no_grad():\n",
    "        for data in loader:\n",
    "            images, labels = data\n",
    "            images, labels = images.to(device), labels.to(device)\n",
    "            outputs_class = net_class(images)\n",
    "            outputs_rej = net_rej(images)\n",
    "            _, predicted = torch.max(outputs_class.data, 1)\n",
    "            _, predicted_rej = torch.max(outputs_rej.data, 1)\n",
    "            batch_size = outputs_class.size()[0]  # batch_size\n",
    "            exp_prediction = expert_fn(images, labels)\n",
    "            for i in range(0, batch_size):\n",
    "                r = (predicted_rej[i].item() == 1)\n",
    "                if r == 0:\n",
    "                    total += 1\n",
    "                    correct += (predicted[i] == labels[i]).item()\n",
    "                    correct_sys += (predicted[i] == labels[i]).item()\n",
    "                if r == 1:\n",
    "                    exp += (exp_prediction[i] == labels[i].item())\n",
    "                    correct_sys += (exp_prediction[i] == labels[i].item())\n",
    "                    exp_total += 1\n",
    "                real_total += 1\n",
    "    cov = str(total) + str(\" out of\") + str(real_total)\n",
    "    to_print = {\"coverage\": cov, \"system accuracy\": 100 * correct_sys / real_total,\n",
    "                \"expert accuracy\": 100 * exp / (exp_total + 0.0002),\n",
    "                \"classifier accuracy\": 100 * correct / (total + 0.0001)}\n",
    "    print(to_print)\n",
    "\n",
    "# data\n",
    "normalize = transforms.Normalize(mean=[x / 255.0 for x in [125.3, 123.0, 113.9]],\n",
    "                                 std=[x / 255.0 for x in [63.0, 62.1, 66.7]])\n",
    "\n",
    "\n",
    "transform_test = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    normalize\n",
    "])\n",
    "\n",
    "if n_dataset == 10:\n",
    "    dataset = 'cifar10'\n",
    "elif n_dataset == 100:\n",
    "    dataset = 'cifar100'\n",
    "\n",
    "kwargs = {'num_workers': 1, 'pin_memory': True}\n",
    "\n",
    "val_loader = torch.utils.data.DataLoader(\n",
    "    datasets.__dict__[dataset.upper()]('../data', train=False, transform=transform_test),\n",
    "    batch_size=128, shuffle=True, **kwargs)\n",
    "\n",
    "\n",
    "class synth_expert:\n",
    "    def __init__(self, k, n_classes):\n",
    "        self.k = k\n",
    "        self.n_classes = n_classes\n",
    "\n",
    "    def predict(self, input, labels):\n",
    "        batch_size = labels.size()[0]  # batch_size\n",
    "        outs = [0] * batch_size\n",
    "        for i in range(0, batch_size):\n",
    "            if labels[i].item() <= self.k:\n",
    "                outs[i] = labels[i].item()\n",
    "            else:\n",
    "                prediction_rand = random.randint(0, self.n_classes - 1)\n",
    "                outs[i] = prediction_rand\n",
    "        return outs\n",
    "\n",
    "expert = synth_expert(k, n_dataset)\n",
    "metrics_print_baseline(model_classifier, model_rejector, expert.predict, n_dataset, val_loader)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Oe0rNcCsOrmK"
   },
   "source": [
    "# Baseline: Confidence\n",
    "Classifier: trained on all data \n",
    "\n",
    "Rejector:\n",
    "\n",
    "1) train model to learn whether expert agrees with target (model_expert)\n",
    "\n",
    "2) Compare confidence of classifier and expert model (model_expert)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "IfUFZEOWOvcX"
   },
   "source": [
    "## Train classifier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "YEWrh7LGOxad"
   },
   "outputs": [],
   "source": [
    "def my_CrossEntropyLoss(outputs, labels):\n",
    "    # m: expert costs, labels: ground truth, n_classes: number of classes\n",
    "    batch_size = outputs.size()[0]  # batch_size\n",
    "    outputs = - torch.log2(outputs[range(batch_size), labels])  # pick the values corresponding to the labels\n",
    "    return torch.sum(outputs) / batch_size\n",
    "\n",
    "def train_classifier(train_loader, model, optimizer, scheduler, epoch, expert_fn, n_classes):\n",
    "    \"\"\"Train for one epoch on the training set\"\"\"\n",
    "    # expertfn: a number here k \n",
    "    batch_time = AverageMeter()\n",
    "    losses = AverageMeter()\n",
    "    top1 = AverageMeter()\n",
    "\n",
    "    # switch to train mode\n",
    "    model.train()\n",
    "\n",
    "    end = time.time()\n",
    "    for i, (input, target) in enumerate(train_loader):\n",
    "        target = target.to(device)\n",
    "        input = input.to(device)\n",
    "\n",
    "        # compute output\n",
    "        output = model(input)\n",
    "        # compute loss\n",
    "        loss = my_CrossEntropyLoss(output, target)\n",
    "\n",
    "        # measure accuracy and record loss\n",
    "        prec1 = accuracy(output.data, target, topk=(1,))[0]\n",
    "        losses.update(loss.data.item(), input.size(0))\n",
    "        top1.update(prec1.item(), input.size(0))\n",
    "\n",
    "        # compute gradient and do SGD step\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        scheduler.step()\n",
    "\n",
    "        # measure elapsed time\n",
    "        batch_time.update(time.time() - end)\n",
    "        end = time.time()\n",
    "\n",
    "        if i % 10 == 0:\n",
    "            print('Epoch: [{0}][{1}/{2}]\\t'\n",
    "                  'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\\t'\n",
    "                  'Loss {loss.val:.4f} ({loss.avg:.4f})\\t'\n",
    "                  'Prec@1 {top1.val:.3f} ({top1.avg:.3f})'.format(\n",
    "                epoch, i, len(train_loader), batch_time=batch_time,\n",
    "                loss=losses, top1=top1))\n",
    "\n",
    "\n",
    "def validate_classifier(val_loader, model, epoch, expert_fn, n_classes):\n",
    "    \"\"\"Perform validation on the validation set\"\"\"\n",
    "    batch_time = AverageMeter()\n",
    "    losses = AverageMeter()\n",
    "    top1 = AverageMeter()\n",
    "\n",
    "    # switch to evaluate mode\n",
    "    model.eval()\n",
    "\n",
    "    end = time.time()\n",
    "    for i, (input, target) in enumerate(val_loader):\n",
    "        target = target.to(device)\n",
    "        input = input.to(device)\n",
    "\n",
    "        # compute output\n",
    "        with torch.no_grad():\n",
    "            output = model(input)\n",
    "        # compute loss\n",
    "        loss = my_CrossEntropyLoss(output, target)\n",
    "\n",
    "        # measure accuracy and record loss\n",
    "        prec1 = accuracy(output.data, target, topk=(1,))[0]\n",
    "        losses.update(loss.data.item(), input.size(0))\n",
    "        top1.update(prec1.item(), input.size(0))\n",
    "\n",
    "        # measure elapsed time\n",
    "        batch_time.update(time.time() - end)\n",
    "        end = time.time()\n",
    "\n",
    "        if i % 10 == 0:\n",
    "            print('Test: [{0}/{1}]\\t'\n",
    "                  'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\\t'\n",
    "                  'Loss {loss.val:.4f} ({loss.avg:.4f})\\t'\n",
    "                  'Prec@1 {top1.val:.3f} ({top1.avg:.3f})'.format(\n",
    "                i, len(val_loader), batch_time=batch_time, loss=losses,\n",
    "                top1=top1))\n",
    "\n",
    "    print(' * Prec@1 {top1.avg:.3f}'.format(top1=top1))\n",
    "\n",
    "    return top1.avg\n",
    "def run_classifier(model, data_aug, n_dataset, expert_fn, epochs):\n",
    "    global best_prec1\n",
    "    # Data loading code\n",
    "    normalize = transforms.Normalize(mean=[x / 255.0 for x in [125.3, 123.0, 113.9]],\n",
    "                                     std=[x / 255.0 for x in [63.0, 62.1, 66.7]])\n",
    "\n",
    "    if data_aug:\n",
    "        transform_train = transforms.Compose([\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Lambda(lambda x: F.pad(x.unsqueeze(0),\n",
    "                                              (4, 4, 4, 4), mode='reflect').squeeze()),\n",
    "            transforms.ToPILImage(),\n",
    "            transforms.RandomCrop(32),\n",
    "            transforms.RandomHorizontalFlip(),\n",
    "            transforms.ToTensor(),\n",
    "            normalize,\n",
    "        ])\n",
    "    else:\n",
    "        transform_train = transforms.Compose([\n",
    "            transforms.ToTensor(),\n",
    "            normalize,\n",
    "        ])\n",
    "    transform_test = transforms.Compose([\n",
    "        transforms.ToTensor(),\n",
    "        normalize\n",
    "    ])\n",
    "\n",
    "    if n_dataset == 10:\n",
    "        dataset = 'cifar10'\n",
    "    elif n_dataset == 100:\n",
    "        dataset = 'cifar100'\n",
    "\n",
    "    kwargs = {'num_workers': 0, 'pin_memory': True}\n",
    "    train_dataset_all = datasets.__dict__[dataset.upper()]('../data', train=True, download=True,\n",
    "                                                           transform=transform_train)\n",
    "    train_size = int(0.90 * len(train_dataset_all))\n",
    "    test_size = len(train_dataset_all) - train_size\n",
    "\n",
    "    train_dataset, test_dataset = torch.utils.data.random_split(train_dataset_all, [train_size, test_size])\n",
    "    train_loader = torch.utils.data.DataLoader(train_dataset,\n",
    "                                               batch_size=128, shuffle=True, **kwargs)\n",
    "    test_loader = torch.utils.data.DataLoader(test_dataset,\n",
    "                                               batch_size=128, shuffle=True, **kwargs)\n",
    "    # get the number of model parameters\n",
    "    print('Number of model parameters: {}'.format(\n",
    "        sum([p.data.nelement() for p in model.parameters()])))\n",
    "\n",
    "    # for training on multiple GPUs.\n",
    "    # Use CUDA_VISIBLE_DEVICES=0,1 to specify which GPUs to use\n",
    "    # model = torch.nn.DataParallel(model).cuda()\n",
    "    model = model.to(device)\n",
    "\n",
    "    # optionally resume from a checkpoint\n",
    "\n",
    "    cudnn.benchmark = True\n",
    "\n",
    "    # define loss function (criterion) and optimizer\n",
    "    optimizer = torch.optim.SGD(model.parameters(), 0.1,\n",
    "                                momentum=0.9, nesterov=True,\n",
    "                                weight_decay=5e-4)\n",
    "\n",
    "    # cosine learning rate\n",
    "    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, len(train_loader) * 200)\n",
    "\n",
    "    for epoch in range(0, epochs):\n",
    "        # train for one epoch\n",
    "        train_classifier(train_loader, model, optimizer, scheduler, epoch, expert_fn, n_dataset)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Q7tgeTGCO9M6"
   },
   "outputs": [],
   "source": [
    "n_dataset = 10  # cifar-10\n",
    "model_classifier = WideResNet(28, n_dataset, 4, dropRate=0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "7edX9eJYOtys"
   },
   "outputs": [],
   "source": [
    "run_classifier(model_classifier, True, n_dataset, 1, 200)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Ce565cfCPLn3"
   },
   "source": [
    "## Model expert confidence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "selqwP7MPNiI"
   },
   "outputs": [],
   "source": [
    "def my_CrossEntropyLoss(outputs, labels):\n",
    "    # m: expert costs, labels: ground truth, n_classes: number of classes\n",
    "    batch_size = outputs.size()[0]  # batch_size\n",
    "    outputs = - torch.log2(outputs[range(batch_size), labels])  # pick the values corresponding to the labels\n",
    "    return torch.sum(outputs) / batch_size\n",
    "\n",
    "def train_expert(train_loader, model, optimizer, scheduler, epoch, expert_fn, n_classes):\n",
    "    \"\"\"Train for one epoch on the training set\"\"\"\n",
    "    batch_time = AverageMeter()\n",
    "    losses = AverageMeter()\n",
    "    top1 = AverageMeter()\n",
    "\n",
    "    # switch to train mode\n",
    "    model.train()\n",
    "\n",
    "    end = time.time()\n",
    "    for i, (input, target) in enumerate(train_loader):\n",
    "        target = target.to(device)\n",
    "        input = input.to(device)\n",
    "\n",
    "        # compute output\n",
    "        output = model(input)\n",
    "        # compute new target\n",
    "        batch_size = output.size()[0]  # batch_size\n",
    "        m = expert_fn(input, target)\n",
    "        for j in range(0, batch_size):\n",
    "            m[j] = 1 - (m[j] == target[j].item())\n",
    "        m = torch.tensor(m)\n",
    "        m = m.to(device)\n",
    "        # compute loss\n",
    "        loss = my_CrossEntropyLoss(output, m)\n",
    "\n",
    "        # measure accuracy and record loss\n",
    "        prec1 = accuracy(output.data, m, topk=(1,))[0]\n",
    "        losses.update(loss.data.item(), input.size(0))\n",
    "        top1.update(prec1.item(), input.size(0))\n",
    "\n",
    "        # compute gradient and do SGD step\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        scheduler.step()\n",
    "\n",
    "        # measure elapsed time\n",
    "        batch_time.update(time.time() - end)\n",
    "        end = time.time()\n",
    "\n",
    "        if i % 10 == 0:\n",
    "            print('Epoch: [{0}][{1}/{2}]\\t'\n",
    "                  'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\\t'\n",
    "                  'Loss {loss.val:.4f} ({loss.avg:.4f})\\t'\n",
    "                  'Prec@1 {top1.val:.3f} ({top1.avg:.3f})'.format(\n",
    "                epoch, i, len(train_loader), batch_time=batch_time,\n",
    "                loss=losses, top1=top1))\n",
    "\n",
    "\n",
    "def validate_expert(val_loader, model, epoch, expert_fn, n_classes):\n",
    "    \"\"\"Perform validation on the validation set\"\"\"\n",
    "    batch_time = AverageMeter()\n",
    "    losses = AverageMeter()\n",
    "    top1 = AverageMeter()\n",
    "\n",
    "    # switch to evaluate mode\n",
    "    model.eval()\n",
    "\n",
    "    end = time.time()\n",
    "    for i, (input, target) in enumerate(val_loader):\n",
    "        target = target.to(device)\n",
    "        input = input.to(device)\n",
    "\n",
    "        # compute output\n",
    "        with torch.no_grad():\n",
    "            output = model(input)\n",
    "        # expert prediction\n",
    "        batch_size = output.size()[0]  # batch_size\n",
    "        m = expert_fn(input, target)\n",
    "        for j in range(0, batch_size):\n",
    "            m[j] = 1 - (m[j] == target[j].item())\n",
    "        m = torch.tensor(m)\n",
    "        m = m.to(device)\n",
    "        # compute loss\n",
    "        loss = my_CrossEntropyLoss(output, m)\n",
    "\n",
    "        # measure accuracy and record loss\n",
    "        prec1 = accuracy(output.data, m, topk=(1,))[0]\n",
    "        losses.update(loss.data.item(), input.size(0))\n",
    "        top1.update(prec1.item(), input.size(0))\n",
    "\n",
    "        # measure elapsed time\n",
    "        batch_time.update(time.time() - end)\n",
    "        end = time.time()\n",
    "\n",
    "        if i % 10 == 0:\n",
    "            print('Test: [{0}/{1}]\\t'\n",
    "                  'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\\t'\n",
    "                  'Loss {loss.val:.4f} ({loss.avg:.4f})\\t'\n",
    "                  'Prec@1 {top1.val:.3f} ({top1.avg:.3f})'.format(\n",
    "                i, len(val_loader), batch_time=batch_time, loss=losses,\n",
    "                top1=top1))\n",
    "\n",
    "    print(' * Prec@1 {top1.avg:.3f}'.format(top1=top1))\n",
    "\n",
    "    return top1.avg\n",
    "best_prec1 = 0\n",
    "def run_expert(model, data_aug, n_dataset, expert_fn, epochs):\n",
    "    global best_prec1\n",
    "    # Data loading code\n",
    "    normalize = transforms.Normalize(mean=[x / 255.0 for x in [125.3, 123.0, 113.9]],\n",
    "                                     std=[x / 255.0 for x in [63.0, 62.1, 66.7]])\n",
    "\n",
    "    if data_aug:\n",
    "        transform_train = transforms.Compose([\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Lambda(lambda x: F.pad(x.unsqueeze(0),\n",
    "                                              (4, 4, 4, 4), mode='reflect').squeeze()),\n",
    "            transforms.ToPILImage(),\n",
    "            transforms.RandomCrop(32),\n",
    "            transforms.RandomHorizontalFlip(),\n",
    "            transforms.ToTensor(),\n",
    "            normalize,\n",
    "        ])\n",
    "    else:\n",
    "        transform_train = transforms.Compose([\n",
    "            transforms.ToTensor(),\n",
    "            normalize,\n",
    "        ])\n",
    "    transform_test = transforms.Compose([\n",
    "        transforms.ToTensor(),\n",
    "        normalize\n",
    "    ])\n",
    "\n",
    "    if n_dataset == 10:\n",
    "        dataset = 'cifar10'\n",
    "    elif n_dataset == 100:\n",
    "        dataset = 'cifar100'\n",
    "\n",
    "    kwargs = {'num_workers': 0, 'pin_memory': True}\n",
    "    train_dataset_all = datasets.__dict__[dataset.upper()]('../data', train=True, download=True,\n",
    "                                                           transform=transform_train)\n",
    "    train_size = int(0.90 * len(train_dataset_all))\n",
    "    test_size = len(train_dataset_all) - train_size\n",
    "\n",
    "    train_dataset, test_dataset = torch.utils.data.random_split(train_dataset_all, [train_size, test_size])\n",
    "    train_loader = torch.utils.data.DataLoader(train_dataset,\n",
    "                                               batch_size=128, shuffle=True, **kwargs)\n",
    "    test_loader = torch.utils.data.DataLoader(test_dataset,\n",
    "                                               batch_size=128, shuffle=True, **kwargs)\n",
    "\n",
    "    # get the number of model parameters\n",
    "    print('Number of model parameters: {}'.format(\n",
    "        sum([p.data.nelement() for p in model.parameters()])))\n",
    "\n",
    "    # for training on multiple GPUs.\n",
    "    # Use CUDA_VISIBLE_DEVICES=0,1 to specify which GPUs to use\n",
    "    # model = torch.nn.DataParallel(model).cuda()\n",
    "    model = model.to(device)\n",
    "\n",
    "    # optionally resume from a checkpoint\n",
    "\n",
    "    cudnn.benchmark = True\n",
    "\n",
    "    # define loss function (criterion) and optimizer\n",
    "    optimizer = torch.optim.SGD(model.parameters(), 0.1,\n",
    "                                momentum=0.9, nesterov=True,\n",
    "                                weight_decay=5e-4)\n",
    "\n",
    "    # cosine learning rate\n",
    "    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, len(train_loader) * 200)\n",
    "\n",
    "    for epoch in range(0, epochs):\n",
    "        # train for one epoch\n",
    "        train_expert(train_loader, model, optimizer, scheduler, epoch, expert_fn, n_dataset)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "hnQY-wjhum01"
   },
   "outputs": [],
   "source": [
    "k = 5 # max k is 10 and min k is 0, irrelevant here\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "v0UdQkAFReoi"
   },
   "outputs": [],
   "source": [
    "n_dataset = 10  # cifar-10\n",
    "\n",
    "\n",
    "class synth_expert:\n",
    "    def __init__(self, k, n_classes):\n",
    "        self.k = k\n",
    "        self.n_classes = n_classes\n",
    "\n",
    "    def predict(self, input, labels):\n",
    "        batch_size = labels.size()[0]  # batch_size\n",
    "        outs = [0] * batch_size\n",
    "        for i in range(0, batch_size):\n",
    "            if labels[i].item() <= self.k:\n",
    "                outs[i] = labels[i].item()\n",
    "            else:\n",
    "                prediction_rand = random.randint(0, self.n_classes - 1)\n",
    "                outs[i] = prediction_rand\n",
    "        return outs\n",
    "\n",
    "\n",
    "model_expert = WideResNet(10, 2, 4, dropRate=0)\n",
    "expert = synth_expert(k, n_dataset)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "colab_type": "code",
    "id": "3ng7oN3qWApQ",
    "outputId": "6185716e-24a8-4c9f-9bc6-f7206c55e3d4"
   },
   "outputs": [],
   "source": [
    "run_expert(model_expert, True, n_dataset, expert.predict, 160)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluate\n",
    "- coverage: percentage of examples where classifier predicts\n",
    "- system accuracy: accuracy of combined system on all data (with deferral)\n",
    "- classifier accuracy: accuracy of classifier on non-deferred examples\n",
    "- expert accuracy: accuracy of expert on deferred examples\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 92
    },
    "colab_type": "code",
    "id": "BFceCtniQmk9",
    "outputId": "a8006283-e58c-43a0-ac94-2e08b8cd645d"
   },
   "outputs": [],
   "source": [
    "def metrics_print_2step(net_mod, net_exp, expert_fn, n_classes, loader):\n",
    "    correct = 0\n",
    "    correct_sys = 0\n",
    "    exp = 0\n",
    "    exp_total = 0\n",
    "    total = 0\n",
    "    real_total = 0\n",
    "    with torch.no_grad():\n",
    "        for data in loader:\n",
    "            images, labels = data\n",
    "            images, labels = images.to(device), labels.to(device)\n",
    "            outputs_mod = net_mod(images)\n",
    "            outputs_exp = net_exp(images)\n",
    "            _, predicted = torch.max(outputs_mod.data, 1)\n",
    "            _, predicted_exp = torch.max(outputs_exp.data, 1)\n",
    "            batch_size = outputs_mod.size()[0]  # batch_size\n",
    "            exp_prediction = expert_fn(images, labels)\n",
    "            for i in range(0, batch_size):\n",
    "                r_score = 1 - outputs_mod.data[i][predicted[i].item()].item()\n",
    "                r_score = r_score - outputs_exp.data[i][1].item()\n",
    "                r = 0\n",
    "                if r_score >= 0:\n",
    "                    r = 1\n",
    "                else:\n",
    "                    r = 0\n",
    "                if r == 0:\n",
    "                    total += 1\n",
    "                    correct += (predicted[i] == labels[i]).item()\n",
    "                    correct_sys += (predicted[i] == labels[i]).item()\n",
    "                if r == 1:\n",
    "                    exp += (exp_prediction[i] == labels[i].item())\n",
    "                    correct_sys += (exp_prediction[i] == labels[i].item())\n",
    "                    exp_total += 1\n",
    "                real_total += 1\n",
    "    cov = str(total) + str(\" out of\") + str(real_total)\n",
    "    to_print = {\"coverage\": cov, \"system accuracy\": 100 * correct_sys / real_total,\n",
    "                \"expert accuracy\": 100 * exp / (exp_total + 0.0002),\n",
    "                \"classifier accuracy\": 100 * correct / (total + 0.0001)}\n",
    "    print(to_print)\n",
    "    \n",
    "normalize = transforms.Normalize(mean=[x / 255.0 for x in [125.3, 123.0, 113.9]],\n",
    "                                 std=[x / 255.0 for x in [63.0, 62.1, 66.7]])\n",
    "\n",
    "transform_test = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    normalize\n",
    "])\n",
    "\n",
    "if n_dataset == 10:\n",
    "    dataset = 'cifar10'\n",
    "elif n_dataset == 100:\n",
    "    dataset = 'cifar100'\n",
    "\n",
    "kwargs = {'num_workers': 1, 'pin_memory': True}\n",
    "\n",
    "val_loader = torch.utils.data.DataLoader(\n",
    "    datasets.__dict__[dataset.upper()]('../data', train=False, transform=transform_test),\n",
    "    batch_size=128, shuffle=True, **kwargs)\n",
    "\n",
    "metrics_print_2step(model_classifier, model_expert, expert.predict, n_dataset, val_loader)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "-MlRVWnvPjGA"
   },
   "source": [
    "# Baseline: Oracle\n",
    "Implements oracle behavior with knowledge of label: defer if label is in the first k classes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 92
    },
    "colab_type": "code",
    "id": "YhYxW4iBRtct",
    "outputId": "7b4178da-1b97-4a3f-9e31-3e2213b7e3b4"
   },
   "outputs": [],
   "source": [
    "def metrics_print_oracle(net_mod, expert_fn, n_classes, loader, k):\n",
    "    correct = 0\n",
    "    correct_sys = 0\n",
    "    exp = 0\n",
    "    exp_total = 0\n",
    "    total = 0\n",
    "    real_total = 0\n",
    "    with torch.no_grad():\n",
    "        for data in loader:\n",
    "            images, labels = data\n",
    "            images, labels = images.to(device), labels.to(device)\n",
    "            outputs_mod = net_mod(images)\n",
    "            _, predicted = torch.max(outputs_mod.data, 1)\n",
    "            batch_size = outputs_mod.size()[0]  # batch_size\n",
    "            exp_prediction = expert_fn(images, labels)\n",
    "            for i in range(0, batch_size):\n",
    "\n",
    "                r = 0\n",
    "                if (labels[i] <= k).item():\n",
    "                    r = 1\n",
    "                else:\n",
    "                    r = 0\n",
    "                if r == 0:\n",
    "                    total += 1\n",
    "                    correct += (predicted[i] == labels[i]).item()\n",
    "                    correct_sys += (predicted[i] == labels[i]).item()\n",
    "                if r == 1:\n",
    "                    exp += (exp_prediction[i] == labels[i].item())\n",
    "                    correct_sys += (exp_prediction[i] == labels[i].item())\n",
    "                    exp_total += 1\n",
    "                real_total += 1\n",
    "    cov = str(total) + str(\" out of\") + str(real_total)\n",
    "    to_print = {\"coverage\": cov, \"system accuracy\": 100 * correct_sys / real_total,\n",
    "                \"expert accuracy\": 100 * exp / (exp_total + 0.0002),\n",
    "                \"classifier accuracy\": 100 * correct / (total + 0.0001)}\n",
    "    print(to_print)\n",
    "    \n",
    "normalize = transforms.Normalize(mean=[x / 255.0 for x in [125.3, 123.0, 113.9]],\n",
    "                                 std=[x / 255.0 for x in [63.0, 62.1, 66.7]])\n",
    "\n",
    "\n",
    "class synth_expert:\n",
    "    def __init__(self, k, n_classes):\n",
    "        self.k = k\n",
    "        self.n_classes = n_classes\n",
    "\n",
    "    def predict(self, input, labels):\n",
    "        batch_size = labels.size()[0]  # batch_size\n",
    "        outs = [0] * batch_size\n",
    "        for i in range(0, batch_size):\n",
    "            if labels[i].item() <= self.k:\n",
    "                coin = np.random.binomial(1, 1, 1)[0] # probability of coin makes expert noisy\n",
    "                if coin:\n",
    "                    outs[i] = labels[i].item()\n",
    "                else:\n",
    "                    prediction_rand = random.randint(0, self.n_classes - 1)\n",
    "                    outs[i] = prediction_rand\n",
    "            else:\n",
    "                prediction_rand = random.randint(0, self.n_classes - 1)\n",
    "                outs[i] = prediction_rand\n",
    "        return outs\n",
    "\n",
    "\n",
    "expert = synth_expert(k, n_dataset)\n",
    "\n",
    "transform_test = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    normalize\n",
    "])\n",
    "\n",
    "if n_dataset == 10:\n",
    "    dataset = 'cifar10'\n",
    "elif n_dataset == 100:\n",
    "    dataset = 'cifar100'\n",
    "\n",
    "kwargs = {'num_workers': 1, 'pin_memory': True}\n",
    "\n",
    "val_loader = torch.utils.data.DataLoader(\n",
    "    datasets.__dict__[dataset.upper()]('../data', train=False, transform=transform_test),\n",
    "    batch_size=128, shuffle=True, **kwargs)\n",
    "\n",
    "metrics_print_oracle(model_classifier, expert.predict, n_dataset, val_loader, k)\n"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "CIFAR-10- baselines.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}