{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Benchmark of attacks.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import torch\nfrom tqdm import tqdm\n\nimport chop\nfrom chop.optim import minimize_frank_wolfe, minimize_pgd, minimize_pgd_madry, minimize_three_split\nfrom chop.data import load_cifar10\nfrom chop.adversary import Adversary\n\nfrom robustbench.utils import load_model\n\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\n\nbatch_size = 50\nn_examples = 10000\nloaders = load_cifar10(test_batch_size=batch_size, data_dir='~/datasets')\nloader = loaders.test\n\nmodel_name = 'Engstrom2019Robustness'\nmodel = load_model(model_name, norm='Linf').to(device)\ncriterion = torch.nn.CrossEntropyLoss(reduction='none')\n\n# Define the perturbation constraint set\nmax_iter = 20\nalpha = 8 / 255.\nconstraint = chop.constraints.LinfBall(alpha)\n\n\nprint(f\"Evaluating model {model_name} on L{constraint.p} ball({alpha}).\")\n\nn_correct = 0\nn_correct_adv_pgd_madry = 0\nn_correct_adv_pgd = 0\nn_correct_adv_split = 0\nn_correct_adv_fw = 0\n\nadversary_pgd = Adversary(minimize_pgd)\nadversary_pgd_madry = Adversary(minimize_pgd_madry)\nadversary_split = Adversary(minimize_three_split)\nadversary_fw = Adversary(minimize_frank_wolfe)\n\nfor k, (data, target) in tqdm(enumerate(loader), total=len(loader)):\n    data = data.to(device)\n    target = target.to(device)\n\n    def image_constraint_prox(delta, step_size=None):\n        \"\"\"Projects perturbation delta\n        so that 0. <= data + delta <= 1.\"\"\"\n\n        adv_img = torch.clamp(data + delta, 0, 1)\n        delta = adv_img - data\n        return delta\n\n    def prox(delta, step_size=None):\n        delta = constraint.prox(delta, step_size)\n        delta = image_constraint_prox(delta, step_size)\n        return delta\n\n\n    _, delta_pgd = adversary_pgd.perturb(data, target, model, criterion,\n                                         use_best=True,\n                                         step='backtracking',\n                                         prox=prox,\n                                         max_iter=max_iter)\n\n    delta_pgd_madry = torch.zeros_like(data)\n    # _, delta_pgd_madry = adversary_pgd_madry.perturb(data, target, model,\n    #                                                  criterion,\n    #                                                  use_best=False,\n    #                                                  prox=prox,\n    #                                                  lmo=constraint.lmo,\n    #                                                  step=2. / max_iter,\n    #                                                  max_iter=max_iter)\n\n    delta_split = torch.zeros_like(data)\n    # _, delta_split = adversary_split.perturb(data, target, model,\n    #                                          criterion,\n    #                                          use_best=False,\n    #                                          prox1=constraint.prox,\n    #                                          prox2=image_constraint_prox,\n    #                                          max_iter=max_iter)\n\n    delta_fw = torch.zeros_like(data)\n    # _, delta_fw = adversary_fw.perturb(data, target, model, criterion,\n    #                                    lmo=constraint.lmo,\n    #                                    step='sublinear',\n    #                                    max_iter=max_iter\n    #                                    )\n\n    label = torch.argmax(model(data), dim=-1)\n    n_correct += (label == target).sum().item()\n\n    adv_label_pgd_madry = torch.argmax(model(data + delta_pgd_madry), dim=-1)\n    n_correct_adv_pgd_madry += (adv_label_pgd_madry == target).sum().item()\n\n    adv_label_pgd = torch.argmax(model(data + delta_pgd), dim=-1)\n    n_correct_adv_pgd += (adv_label_pgd == target).sum().item()\n\n    adv_label_split = torch.argmax(model(data + delta_split), dim=-1)\n    n_correct_adv_split += (adv_label_split == target).sum().item()\n\n    adv_label_fw = torch.argmax(model(data + delta_fw), dim=-1)\n    n_correct_adv_fw += (adv_label_fw == target).sum().item()\n\n\naccuracy = n_correct / n_examples\naccuracy_adv_pgd_madry = n_correct_adv_pgd_madry / n_examples\naccuracy_adv_pgd = n_correct_adv_pgd / n_examples\naccuracy_adv_split = n_correct_adv_split / n_examples\naccuracy_adv_fw = n_correct_adv_fw / n_examples\n\nprint(f\"Accuracy: {accuracy:.4f}\")\nprint(f\"RobustAccuracy PGD Madry: {accuracy_adv_pgd_madry:.4f}\")\nprint(f\"RobustAccuracy PGD: {accuracy_adv_pgd:.4f}\")\nprint(f\"RobustAccuracy Splitting: {accuracy_adv_split:.4f}\")\nprint(f\"RobustAccuracy FW: {accuracy_adv_fw:.4f}\")"
      ]
    }
  ],
  "metadata": {
    "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.8.5"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}