{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Example of robust training on CIFAR10.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from chop.adversary import Adversary\nimport torch\nfrom tqdm import tqdm\n\nimport chop\n\nfrom torch.optim import SGD\n\nfrom torchvision import models\n\ndevice = torch.device('cuda' if torch.cuda.is_available()\n                      else 'cpu')\n\nn_epochs = 128\nbatch_size = 100\n\nloaders = chop.data.load_cifar10(train_batch_size=batch_size,\n                                 test_batch_size=batch_size,\n                                 data_dir='~/datasets')\ntrainloader, testloader = loaders.train, loaders.test\n\nmodel = models.resnet18(pretrained=False)\nmodel.to(device)\n\ncriterion = torch.nn.CrossEntropyLoss()\n\n# TODO: use learning rate schedulers\noptimizer = SGD(model.parameters(), lr=.1, momentum=.9, weight_decay=5e-4)\n\n# Define the perturbation constraint set\nmax_iter_train = 7\nmax_iter_test = 20\nalpha = 8. / 255\nconstraint = chop.constraints.LinfBall(alpha)\ncriterion_adv = torch.nn.CrossEntropyLoss(reduction='none')\n\nprint(f\"Training on L{constraint.p} ball({alpha}).\")\n\n\nadversary = Adversary(chop.optim.minimize_pgd_madry)\n\nfor _ in range(n_epochs):\n\n    # Train\n    n_correct = 0\n    n_correct_adv = 0\n\n    model.train()\n\n    for k, (data, target) in tqdm(enumerate(trainloader), total=len(trainloader)):\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        _, delta = adversary.perturb(data, target, model,\n                                     criterion_adv,\n                                     prox=prox,\n                                     lmo=constraint.lmo,\n                                     step=2. / max_iter_train,\n                                     max_iter=max_iter_train)\n\n        output = model(data)\n        output_adv = model(data + delta)\n        loss = criterion(output_adv, target)\n        loss.backward()\n\n        pred = torch.argmax(output, dim=-1)\n        pred_adv = torch.argmax(output_adv, dim=-1)\n\n        n_correct += (pred == target).sum().item()\n        n_correct_adv += (pred_adv == target).sum().item()\n\n    print(f\"Train Accuracy: {n_correct / 50000.}\")\n    print(f\"Train Adv Accuracy: {n_correct_adv / 50000.}\")\n\n    # Test\n    n_correct = 0\n    n_correct_adv = 0\n\n    model.eval()\n\n    for k, (data, target) in tqdm(enumerate(testloader), total=len(testloader)):\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        _, delta = adversary.perturb(data, target, model,\n                                     criterion_adv,\n                                     prox=prox,\n                                     lmo=constraint.lmo,\n                                     step=2. / max_iter_test,\n                                     max_iter=max_iter_test)\n\n        with torch.no_grad():\n            output = model(data)\n            output_adv = model(data + delta)\n\n            pred = torch.argmax(output, dim=-1)\n            pred_adv = torch.argmax(output_adv, dim=-1)\n\n            n_correct += (pred == target).sum().item()\n            n_correct_adv += (pred_adv == target).sum().item()\n\n    print(f\"Test Accuracy: {n_correct / 10000.}\")\n    print(f\"Test Adv Accuracy: {n_correct_adv / 10000.}\")"
      ]
    }
  ],
  "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
}