{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Visualizing Adversarial Examples\n\nThis example shows how to generate and plot adversarial examples for a batch of datapoints from CIFAR-10,\nand compares the examples from different constraint sets, penalizations and solvers.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from itertools import product\n\nimport numpy as np\n\nimport torch\nimport torchvision\n\nfrom robustbench.data import load_cifar10\nfrom robustbench.utils import load_model\n\nimport matplotlib.pyplot as plt\n\nimport chop\nfrom chop.image import group_patches, matplotlib_imshow_batch\nfrom chop.logging import Trace\n\nfrom sklearn.metrics import f1_score\n\nfrom torch.autograd import profiler\n\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n# device = torch.device('cpu')\n\nbatch_size = 7\n\n# Note that this example uses load_cifar10 from the robustbench library\ndata, target = load_cifar10(n_examples=batch_size, data_dir='~/datasets')\ndata = data.to(device)\ntarget = target.to(device)\n\nclasses = ('plane', 'car', 'bird', 'cat',\n           'deer', 'dog', 'frog', 'horse',\n           'ship', 'truck')\n\n\nmodel = load_model('Standard')  # Can be changed to any model from the robustbench model zoo\nmodel = model.to(device)\ncriterion = torch.nn.CrossEntropyLoss(reduction='none')\n\n# Define the constraint set + initial point\nprint(\"L2 norm constraint.\")\nalpha = 127 / 255\nconstraint = chop.constraints.L2Ball(alpha)\n\n\ndef image_constraint_prox(delta, step_size=None):\n    adv_img = torch.clamp(data + delta, 0, 1)\n    delta = adv_img - data\n    return delta\n\n\ndef 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\nadversary = chop.Adversary(chop.optim.minimize_pgd_madry)\ncallback_L2 = Trace()\n_, delta = adversary.perturb(data, target, model, criterion,\n                             prox=prox,\n                             lmo=constraint.lmo,\n                             max_iter=20,\n                             step=2. / 20,\n                             callback=callback_L2)\n\n# Plot adversarial images\nfig, ax = plt.subplots(nrows=7, ncols=batch_size, figsize=(16, 14))\n\n# Plot clean data\nmatplotlib_imshow_batch(data, labels=[classes[k] for k in target], axes=ax[0, :],\n                        title=\"Original images\")\n\n# Adversarial Lp images\nadv_output = model(data + delta)\nadv_labels = torch.argmax(adv_output, dim=-1)\nmatplotlib_imshow_batch(data + delta, labels=[classes[k] for k in adv_labels], axes=ax[1, :],\n                        title=f'L{constraint.p}')\n\n# Perturbation\nmatplotlib_imshow_batch(abs(delta), axes=ax[4, :], normalize=True,\n                        title=f'L{constraint.p}')\n\n\nprint(\"GroupL1 constraint.\")\n\ngroups = group_patches(x_patch_size=8, y_patch_size=8)\n\nfor eps in [5e-2]:\n    alpha = eps * len(groups)\n    constraint_group = chop.constraints.GroupL1Ball(alpha, groups)\n    adversary_group = chop.Adversary(chop.optim.minimize_frank_wolfe)\n\n    # callback_group = Trace(callable=lambda kw: criterion(model(data + kw['x']), target))\n    callback_group = Trace()\n\n    with profiler.profile() as prof:\n        _, delta_group = adversary_group.perturb(data, target, model, criterion,\n                                                lmo=constraint_group.lmo,\n                                                max_iter=20,\n                                                callback=callback_group)\n\n    \n    delta_group = image_constraint_prox(delta_group)\n\n    # Show adversarial examples and perturbations\n    adv_output_group = model(data + delta_group)\n    adv_labels_group = torch.argmax(adv_output_group, dim=-1)\n\n    matplotlib_imshow_batch(data + delta_group, labels=(classes[k] for k in adv_labels_group),\n                            axes=ax[2, :],\n                            title='Group Lasso')\n\n    matplotlib_imshow_batch(abs(delta_group), axes=ax[5, :], normalize=True,\n                            title='Group Lasso')\n\n    print(f\"F1 score: {f1_score(target.detach().cpu(), adv_labels_group.detach().cpu(), average='macro'):.3f}\"\n          f\" for alpha={alpha:.4f}\")\n\nprint(\"Nuclear norm ball adv examples\")\n\nfor alpha in [.5]:\n    constraint_nuc = chop.constraints.NuclearNormBall(alpha)\n\n\n    def prox_nuc(delta, step_size=None):\n        delta = constraint_nuc.prox(delta, step_size)\n        delta = image_constraint_prox(delta, step_size)\n        return delta\n\n\n    adversary = chop.Adversary(chop.optim.minimize_frank_wolfe)\n    callback_nuc = Trace()\n\n    _, delta_nuc = adversary.perturb(data, target, model, criterion,\n                                #  prox=prox,\n                                lmo=constraint_nuc.lmo,\n                                max_iter=20,\n                                #  step=2. / 20,\n                                callback=callback_nuc)\n\n    # Clamp last iterate to image space\n    delta_nuc = image_constraint_prox(delta_nuc)\n\n    # Add nuclear examples to plot\n    adv_output_nuc = model(data + delta_nuc)\n    adv_labels_nuc = torch.argmax(adv_output_nuc, dim=-1)\n\n    matplotlib_imshow_batch(data + delta_nuc, labels=(classes[k] for k in adv_labels_nuc),\n                            axes=ax[3, :],\n                            title='Nuclear Norm')\n\n    matplotlib_imshow_batch(abs(delta_nuc), axes=ax[6, :], normalize=True,\n                            title='Nuclear Norm')\n\n    print(f\"F1 score: {f1_score(target.detach().cpu(), adv_labels_nuc.detach().cpu(), average='macro'):.3f}\"\n          f\" for alpha={alpha:.4f}\")\n\nplt.tight_layout()\nplt.show()\n\n\n# TODO refactor this in functions\n\n# Plot group lasso loss values\nfig, ax = plt.subplots(figsize=(6, 10), nrows=batch_size, sharex=True)\nfor k in range(batch_size):\n    ax[k].plot([-trace[k] for trace in callback_group.trace_f])\nplt.tight_layout()\nplt.show()\n\n# Plot loss functions per datapoint\nfig, ax = plt.subplots(figsize=(6, 10), nrows=batch_size, sharex=True)\nfor k in range(batch_size):\n    ax[k].plot([-trace[k] for trace in callback_nuc.trace_f])\n\nplt.tight_layout()\nplt.show()\n\n# Plot loss functions per datapoint\nfig, ax = plt.subplots(figsize=(6, 10), nrows=batch_size, sharex=True)\nfor k in range(batch_size):\n    ax[k].plot([-trace[k] for trace in callback_L2.trace_f])\n\nplt.tight_layout()\nplt.show()"
      ]
    }
  ],
  "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
}