{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Constrained neural network training.\nTrains a LeNet5 model on MNIST using constraints on the weights.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from tqdm import tqdm\n\nimport numpy as np\nimport torch\nfrom torch import nn\n\nfrom easydict import EasyDict\n\nfrom advertorch.test_utils import LeNet5\nfrom advertorch_examples.utils import get_mnist_train_loader\nfrom advertorch_examples.utils import get_mnist_test_loader\n\nimport chop\n\n# Setup\ntorch.manual_seed(0)\n\nuse_cuda = torch.cuda.is_available()\ndevice = torch.device(\"cuda\" if use_cuda else \"cpu\")\n\n# Data Loaders\ntrain_loader = get_mnist_train_loader(batch_size=50, shuffle=True)\ntest_loader = get_mnist_test_loader(batch_size=512, shuffle=True)\n\n# Model setup\nmodel = LeNet5()\nmodel.to(device)\n\ncriterion = nn.CrossEntropyLoss()\n\n# Outer optimization parameters\nnb_epochs = 20\nmomentum = .9\nlr = 0.3\n\n# Choose optimizer here\n# optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\n\n# Make constraints\n\nalpha = 1.\nconstraint = chop.constraints.LinfBall(alpha)\n\n# Project model parameters in the constraint set.\nconstraint.make_feasible(model)\n\noptimizer = chop.stochastic.FrankWolfe(model.parameters(), constraint, lr=lr, momentum=momentum)\n\n# Training loop\nfor epoch in range(nb_epochs):\n    model.train()\n    train_loss = 0.\n    for data, target in tqdm(train_loader, desc=f'Training epoch {epoch}/{nb_epochs - 1}'):\n        data, target = data.to(device), target.to(device)\n\n        optimizer.zero_grad()\n        loss = criterion(model(data), target)\n        loss.backward()\n        optimizer.step()\n\n        train_loss += loss.item()\n    train_loss /= len(train_loader)\n    print(f'Training loss: {train_loss:.3f}')\n    # TODO: get accuracy\n\n    # Evaluate on clean and adversarial test data\n\n    model.eval()\n    report = EasyDict(nb_test=0, correct=0, correct_adv_pgd=0,\n                      correct_adv_pgd_madry=0,\n                      correct_adv_fw=0, correct_adv_mfw=0)\n\n    for data, target in tqdm(test_loader, desc=f'Val epoch {epoch}/{nb_epochs - 1}'):\n        data, target = data.to(device), target.to(device)\n\n        # Compute corresponding predictions        \n        _, pred = model(data).max(1)\n\n        # Get clean accuracies\n        report.nb_test += data.size(0)\n        report.correct += pred.eq(target).sum().item()\n\n    print(f'Val acc on clean examples (%): {report.correct / report.nb_test * 100.:.3f}')"
      ]
    }
  ],
  "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
}