Adversarial_defense.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Adversarial defense.ipynb",
      "provenance": [],
      "authorship_tag": "ABX9TyOLClrRw6AY8j9Zm75X7yGk",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/GiuliaLanzillotta/exercises/blob/master/Adversarial_defense.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GvBIVuZkyVgx"
      },
      "source": [
        "# Adversarial defense \n",
        "\n",
        "Today we'll experiment with adversarial training as an adversarial defense technique. <br>\n",
        "\n",
        "More specifically we'll employ [PGD](https://arxiv.org/pdf/1706.06083.pdf) and [TRADES](https://arxiv.org/pdf/1901.08573.pdf) attack to make our net more robust during training. \n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fIuQvCcothZn"
      },
      "source": [
        "import os\n",
        "from tqdm import tqdm\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.optim as optim\n",
        "from torchvision import datasets, transforms"
      ],
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6roEhzyFy2AV"
      },
      "source": [
        "We'll be again be working with the MNIST dataset "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r2i30W5jyudh"
      },
      "source": [
        "  \n",
        "  ## 1. Define a shallow ReLU network "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "H28j_VzkysfM"
      },
      "source": [
        "class Net(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(Net, self).__init__()\n",
        "        self.fc1 = nn.Linear(28 * 28, 200)\n",
        "        self.fc2 = nn.Linear(200, 10)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = x.view((-1, 28 * 28))\n",
        "        x = F.relu(self.fc1(x))\n",
        "        x = self.fc2(x)\n",
        "        return x"
      ],
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "51VDPYKZzICO"
      },
      "source": [
        "Let's also add a normalisation layer.<br>\n",
        "It will be inserted as a first \"layer\" to the network. This allows us to search for adverserial examples to the real image, rather than to the normalized image. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FfMlkWxNzLnj"
      },
      "source": [
        "class Normalize(nn.Module):\n",
        "    def forward(self, x):\n",
        "        return (x - 0.1307) / 0.3081"
      ],
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6Du0ucKDzOsj"
      },
      "source": [
        "And let's set a few hyperparameters"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nGSJ2zjAzM0n"
      },
      "source": [
        "batch_size = 512\n",
        "seed = 42\n",
        "learning_rate = 0.01\n",
        "num_epochs = 10 \n",
        "eps = 0.1 #PGD parameter (defines the magnitude of the perturbation)\n",
        "k = 7 #PGD steps\n",
        "trades_fact = 1.0 #TRADES lambda "
      ],
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AW2fCyeAzzlP"
      },
      "source": [
        "A few more lines of preparatory code ..."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "neFeWovqzv5o",
        "outputId": "a0b4cf87-b5eb-4043-830d-538b832c1364",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        }
      },
      "source": [
        "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
        "torch.manual_seed(seed)"
      ],
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<torch._C.Generator at 0x7fe9c96a9810>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 5
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "h4zvI-820SBK"
      },
      "source": [
        "model = nn.Sequential(Normalize(), Net())\n",
        "model = model.to(device)"
      ],
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eZ60YN2Gz8tl"
      },
      "source": [
        "\n",
        "  ## 2. Load dataset (MNIST)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "B7nzjTPSz7dX"
      },
      "source": [
        "# Warning: running this will download the data locally\n",
        "train_dataset = datasets.MNIST('mnist_data/', train=True, download=True, transform=transforms.Compose([transforms.ToTensor()]))\n",
        "test_dataset = datasets.MNIST('mnist_data/', train=False, download=True, transform=transforms.Compose([transforms.ToTensor()]))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fR88BXlY0awZ"
      },
      "source": [
        "train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n",
        "test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False)"
      ],
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SXZTDJGz0jL5"
      },
      "source": [
        "\n",
        "## 3. Implement the defenses "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qUP7XoNZ0hpt"
      },
      "source": [
        "def fgsm_step(x,y, eps, net):\n",
        "  \"\"\"\n",
        "  Implements an fgsm step.\n",
        "  \"\"\"\n",
        "  input_ = x.clone().detach_()\n",
        "  input_.requires_grad = True\n",
        "  loss = ce_loss(net(input_), torch.tensor([y], dtype=torch.long))\n",
        "  loss.backward()\n",
        "  x_next = x + eps*torch.sign(input_.grad)\n",
        "  return x_next\n",
        "\n",
        "def fgsm_TRADES_step(x,x_next,eps,net):\n",
        "  \"\"\"\n",
        "  Implements fgsm step for TRADES boundary loss optimization. \n",
        "  TRADES boundary loss encourages the network to have a smoother boundary, with \n",
        "  respect to the perturbation regions. \n",
        "  Basically we enforce the predefined perturbation regions to have a coherent \n",
        "  prediction. \n",
        "  \"\"\"\n",
        "  input_ = x_next.clone().detach_()\n",
        "  input_.requires_grad = True\n",
        "  # The loss is equivalent to the difference between the original prediction\n",
        "  # and the adversarial prediction \n",
        "  target = torch.argmax(net(x),dim=1)\n",
        "  logit = torch.nn.LogSoftmax(dim=1)(net(input_))\n",
        "  loss = torch.nn.NLLLoss()\n",
        "  out = loss(logit,target)\n",
        "  out.backward()\n",
        "  x_next = x_next + eps*torch.sign(input_.grad)\n",
        "  return x_next\n",
        "\n",
        "\n",
        "def get_PGD_adversarial_example(x,y, eps, net, k, method=\"PGD\"):\n",
        "  \"\"\" \n",
        "  Returns adversarial example in epsilon infinity ball around x\n",
        "  using untargeted PGD attack. \n",
        "  \"\"\" \n",
        "  eps_step = 2.5*(eps/k)\n",
        "  x_next = x \n",
        "  # set the adversarial flag to signal whether \n",
        "  # we have found an adversarial example\n",
        "  adv = torch.argmax(net(x), dim=1) != y\n",
        "  if adv: return x_next\n",
        "  for i in range(k): \n",
        "    # take an fgsm step \n",
        "    if method==\"PGD\":x_next = fgsm_step(x_next, y, eps_step, net)\n",
        "    elif method==\"TRADES\":x_next = fgsm_TRADES_step(x, x_next, eps_step, net)\n",
        "    else: raise ValueError(\"Method inserted is not valid. Supported methods: PGD, TRADES.\")\n",
        "    # project back to L infinity ball\n",
        "    delta = x - x_next \n",
        "    delta = torch.clamp(delta, min=-1*eps, max=eps)\n",
        "    x_next = x - delta\n",
        "    # check whether we have an adversarial example \n",
        "    adv = torch.argmax(net(x_next), dim=1) != y\n",
        "    if adv: return x_next\n",
        "  return None\n",
        "\n",
        "def get_PGD_Bmax(x_batch, y_batch, eps, net, k, method=\"PGD\"):\n",
        "  \"\"\" \n",
        "  Returns alternative set of points that maximise \n",
        "  the loss of the newtork \n",
        "  \"\"\"\n",
        "  Bmax = []\n",
        "  for (x,y) in zip(x_batch,y_batch):\n",
        "    x_adv = get_PGD_adversarial_example(x,y, eps, net, k, method=method)\n",
        "    if x_adv is None: x_adv = x \n",
        "    Bmax +=[x_adv[0]]\n",
        "  return torch.stack(Bmax)\n",
        "    \n",
        "def compute_adv_accuracy(x_batch, y_batch, eps, k, lam, net, method=\"PGD\"):\n",
        "  \"\"\"\n",
        "  Returns the adversarial accuracy on the given batch. \n",
        "  \"\"\"\n",
        "  # iterate through the batch:\n",
        "  # for the correctly classified examples check whether there's \n",
        "  # an adversarial example \n",
        "  tot_acc = 0\n",
        "  for (x,y) in zip(x_batch,y_batch):\n",
        "    if torch.argmax(net(x), dim=1) == y and get_PGD_adversarial_example(x,y, eps, net, k, method=method) is None: \n",
        "      tot_acc+=1\n",
        "  return tot_acc\n"
      ],
      "execution_count": 37,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TlEGrghj1syk"
      },
      "source": [
        "\n",
        "## 4. Train and evaluate "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rdGBk-UG1uSe"
      },
      "source": [
        "opt = optim.Adam(params=model.parameters(), lr=learning_rate, weight_decay=1e-5)\n",
        "scheduler = optim.lr_scheduler.StepLR(opt, 15)\n",
        "ce_loss = torch.nn.CrossEntropyLoss()\n",
        "kl_loss = torch.nn.KLDivLoss(reduction='batchmean')"
      ],
      "execution_count": 20,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2it_4_uA16BF"
      },
      "source": [
        "defense = \"TRADES\""
      ],
      "execution_count": 21,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JoEPOzct11Hu",
        "outputId": "ff6c1eea-5abf-44b9-8bec-649ff3016aef",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 711
        }
      },
      "source": [
        "\n",
        "for epoch in range(1, num_epochs + 1):\n",
        "    # Training\n",
        "    for batch_idx, (x_batch, y_batch) in enumerate(tqdm(train_loader)):\n",
        "\n",
        "        x_batch, y_batch = x_batch.to(device), y_batch.to(device)\n",
        "        model.train()  \n",
        "\n",
        "        if defense == 'PGD':\n",
        "          # get Bmax \n",
        "          x_batch_max = get_PGD_Bmax(x_batch, y_batch, eps, model, k)\n",
        "          # compute the loss \n",
        "          out = model(x_batch_max)\n",
        "          loss = ce_loss(out, y_batch)\n",
        "\n",
        "        elif defense == 'TRADES':\n",
        "          # get Bmax \n",
        "          x_batch_max = get_PGD_Bmax(x_batch, y_batch, eps, model, k, method=\"TRADES\")\n",
        "          # compute the loss \n",
        "          out = model(x_batch)\n",
        "          out_max = model(x_batch_max)\n",
        "          # boundary loss \n",
        "          b_loss = torch.nn.NLLLoss()\n",
        "          b_out = b_loss(torch.nn.LogSoftmax(dim=1)(out_max),\n",
        "                         torch.argmax(out,dim=1))\n",
        "          loss = ce_loss(out, y_batch) + trades_fact*b_out\n",
        "\n",
        "        elif defense == 'none':\n",
        "            # standard training\n",
        "            out_nat = model(x_batch)\n",
        "            loss = ce_loss(out_nat, y_batch)\n",
        "        \n",
        "        opt.zero_grad()\n",
        "        loss.backward()\n",
        "        opt.step()\n",
        "        \n",
        "    # Testing\n",
        "    model.eval()\n",
        "    tot_test, tot_acc, tot_adv_acc = 0.0, 0.0, 0.0\n",
        "    for batch_idx, (x_batch, y_batch) in enumerate(tqdm(test_loader)):\n",
        "        x_batch, y_batch = x_batch.to(device), y_batch.to(device)\n",
        "        \n",
        "        out = model(x_batch)\n",
        "        pred = torch.max(out, dim=1)[1]\n",
        "        acc = pred.eq(y_batch).sum().item()\n",
        "\n",
        "        \n",
        "        acc_adv = compute_adv_accuracy(x_batch, y_batch, eps, k, trades_fact, model, method=defense)\n",
        "        \n",
        "        tot_acc += acc\n",
        "        tot_adv_acc += acc_adv\n",
        "        tot_test += x_batch.size()[0]\n",
        "    scheduler.step()\n",
        "    print()\n",
        "    print('Epoch %d: Accuracy %.5lf, Adv Accuracy %.5lf' % (epoch, tot_acc / tot_test, tot_adv_acc / tot_test))\n"
      ],
      "execution_count": 38,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "100%|██████████| 118/118 [04:49<00:00,  2.46s/it]\n",
            "100%|██████████| 20/20 [00:57<00:00,  2.88s/it]\n",
            "  0%|          | 0/118 [00:00<?, ?it/s]"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 1: Accuracy 0.95010, Adv Accuracy 0.59670\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "100%|██████████| 118/118 [05:48<00:00,  2.96s/it]\n",
            "100%|██████████| 20/20 [01:00<00:00,  3.05s/it]\n",
            "  0%|          | 0/118 [00:00<?, ?it/s]"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 2: Accuracy 0.96320, Adv Accuracy 0.64830\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "100%|██████████| 118/118 [05:59<00:00,  3.05s/it]\n",
            "100%|██████████| 20/20 [01:02<00:00,  3.12s/it]\n",
            "  0%|          | 0/118 [00:00<?, ?it/s]"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 3: Accuracy 0.96640, Adv Accuracy 0.67410\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "100%|██████████| 118/118 [06:07<00:00,  3.11s/it]\n",
            "100%|██████████| 20/20 [01:03<00:00,  3.17s/it]\n",
            "  0%|          | 0/118 [00:00<?, ?it/s]"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 4: Accuracy 0.97060, Adv Accuracy 0.70140\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "100%|██████████| 118/118 [06:12<00:00,  3.15s/it]\n",
            "100%|██████████| 20/20 [01:01<00:00,  3.08s/it]\n",
            "  0%|          | 0/118 [00:00<?, ?it/s]"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 5: Accuracy 0.97020, Adv Accuracy 0.66890\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "100%|██████████| 118/118 [06:11<00:00,  3.15s/it]\n",
            "100%|██████████| 20/20 [01:04<00:00,  3.23s/it]\n",
            "  0%|          | 0/118 [00:00<?, ?it/s]"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 6: Accuracy 0.97340, Adv Accuracy 0.71720\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "100%|██████████| 118/118 [06:18<00:00,  3.21s/it]\n",
            "100%|██████████| 20/20 [01:04<00:00,  3.22s/it]\n",
            "  0%|          | 0/118 [00:00<?, ?it/s]"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 7: Accuracy 0.97300, Adv Accuracy 0.72230\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "100%|██████████| 118/118 [06:18<00:00,  3.20s/it]\n",
            "100%|██████████| 20/20 [01:03<00:00,  3.20s/it]\n",
            "  0%|          | 0/118 [00:00<?, ?it/s]"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 8: Accuracy 0.97330, Adv Accuracy 0.71620\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "100%|██████████| 118/118 [06:17<00:00,  3.20s/it]\n",
            "100%|██████████| 20/20 [01:04<00:00,  3.25s/it]\n",
            "  0%|          | 0/118 [00:00<?, ?it/s]"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 9: Accuracy 0.97440, Adv Accuracy 0.73010\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "100%|██████████| 118/118 [06:28<00:00,  3.29s/it]\n",
            "100%|██████████| 20/20 [01:05<00:00,  3.26s/it]"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "\n",
            "Epoch 10: Accuracy 0.97580, Adv Accuracy 0.72910\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "\n"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_RaBY8VJokwT"
      },
      "source": [
        "  Output for PGD attack: \n",
        "    \n",
        "    100%|██████████| 118/118 [03:28<00:00,  1.77s/it]\n",
        "    100%|██████████| 20/20 [00:42<00:00,  2.11s/it]\n",
        "      0%|          | 0/118 [00:00<?, ?it/s]\n",
        "      Epoch 1: Accuracy 0.94200, Adv Accuracy 0.59200\n",
        "    100%|██████████| 118/118 [04:21<00:00,  2.22s/it]\n",
        "    100%|██████████| 20/20 [00:46<00:00,  2.30s/it]\n",
        "      0%|          | 0/118 [00:00<?, ?it/s]\n",
        "      Epoch 2: Accuracy 0.96210, Adv Accuracy 0.67770\n",
        "    100%|██████████| 118/118 [04:28<00:00,  2.27s/it]\n",
        "    100%|██████████| 20/20 [00:46<00:00,  2.30s/it]\n",
        "      0%|          | 0/118 [00:00<?, ?it/s]\n",
        "      Epoch 3: Accuracy 0.96510, Adv Accuracy 0.66470\n",
        "    100%|██████████| 118/118 [04:35<00:00,  2.34s/it]\n",
        "    100%|██████████| 20/20 [00:47<00:00,  2.39s/it]\n",
        "      0%|          | 0/118 [00:00<?, ?it/s]\n",
        "      Epoch 4: Accuracy 0.96780, Adv Accuracy 0.71550\n",
        "    100%|██████████| 118/118 [04:38<00:00,  2.36s/it]\n",
        "    100%|██████████| 20/20 [00:46<00:00,  2.33s/it]\n",
        "      0%|          | 0/118 [00:00<?, ?it/s]\n",
        "      Epoch 5: Accuracy 0.96690, Adv Accuracy 0.69060\n",
        "    100%|██████████| 118/118 [04:39<00:00,  2.37s/it]\n",
        "    100%|██████████| 20/20 [00:47<00:00,  2.38s/it]\n",
        "      0%|          | 0/118 [00:00<?, ?it/s]\n",
        "      Epoch 6: Accuracy 0.96980, Adv Accuracy 0.72700\n",
        "    100%|██████████| 118/118 [04:37<00:00,  2.35s/it]\n",
        "    100%|██████████| 20/20 [00:47<00:00,  2.36s/it]\n",
        "      0%|          | 0/118 [00:00<?, ?it/s]\n",
        "      Epoch 7: Accuracy 0.96770, Adv Accuracy 0.70520\n",
        "    100%|██████████| 118/118 [04:40<00:00,  2.37s/it]\n",
        "    100%|██████████| 20/20 [00:47<00:00,  2.39s/it]\n",
        "      0%|          | 0/118 [00:00<?, ?it/s]\n",
        "      Epoch 8: Accuracy 0.96920, Adv Accuracy 0.72520\n",
        "    100%|██████████| 118/118 [04:39<00:00,  2.37s/it]\n",
        "    100%|██████████| 20/20 [00:47<00:00,  2.38s/it]\n",
        "      0%|          | 0/118 [00:00<?, ?it/s]\n",
        "      Epoch 9: Accuracy 0.96900, Adv Accuracy 0.70910\n",
        "    100%|██████████| 118/118 [04:41<00:00,  2.39s/it]\n",
        "    100%|██████████| 20/20 [00:47<00:00,  2.36s/it]\n",
        "    Epoch 10: Accuracy 0.97080, Adv Accuracy 0.70330\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gX0DQ5UJJwAM"
      },
      "source": [
        "Output for TRADES attack: \n",
        "\n",
        "    100%|██████████| 118/118 [04:49<00:00,  2.46s/it]\n",
        "    100%|██████████| 20/20 [00:57<00:00,  2.88s/it]\n",
        "      0%|          | 0/118 [00:00<?, ?it/s]\n",
        "    Epoch 1: Accuracy 0.95010, Adv Accuracy 0.59670\n",
        "    100%|██████████| 118/118 [05:48<00:00,  2.96s/it]\n",
        "    100%|██████████| 20/20 [01:00<00:00,  3.05s/it]\n",
        "      0%|          | 0/118 [00:00<?, ?it/s]\n",
        "    Epoch 2: Accuracy 0.96320, Adv Accuracy 0.64830\n",
        "    100%|██████████| 118/118 [05:59<00:00,  3.05s/it]\n",
        "    100%|██████████| 20/20 [01:02<00:00,  3.12s/it]\n",
        "      0%|          | 0/118 [00:00<?, ?it/s]\n",
        "    Epoch 3: Accuracy 0.96640, Adv Accuracy 0.67410\n",
        "    100%|██████████| 118/118 [06:07<00:00,  3.11s/it]\n",
        "    100%|██████████| 20/20 [01:03<00:00,  3.17s/it]\n",
        "      0%|          | 0/118 [00:00<?, ?it/s]\n",
        "    Epoch 4: Accuracy 0.97060, Adv Accuracy 0.70140\n",
        "    100%|██████████| 118/118 [06:12<00:00,  3.15s/it]\n",
        "    100%|██████████| 20/20 [01:01<00:00,  3.08s/it]\n",
        "      0%|          | 0/118 [00:00<?, ?it/s]\n",
        "    Epoch 5: Accuracy 0.97020, Adv Accuracy 0.66890\n",
        "    100%|██████████| 118/118 [06:11<00:00,  3.15s/it]\n",
        "    100%|██████████| 20/20 [01:04<00:00,  3.23s/it]\n",
        "      0%|          | 0/118 [00:00<?, ?it/s]\n",
        "    Epoch 6: Accuracy 0.97340, Adv Accuracy 0.71720\n",
        "    100%|██████████| 118/118 [06:18<00:00,  3.21s/it]\n",
        "    100%|██████████| 20/20 [01:04<00:00,  3.22s/it]\n",
        "      0%|          | 0/118 [00:00<?, ?it/s]\n",
        "    Epoch 7: Accuracy 0.97300, Adv Accuracy 0.72230\n",
        "    100%|██████████| 118/118 [06:18<00:00,  3.20s/it]\n",
        "    100%|██████████| 20/20 [01:03<00:00,  3.20s/it]\n",
        "      0%|          | 0/118 [00:00<?, ?it/s]\n",
        "    Epoch 8: Accuracy 0.97330, Adv Accuracy 0.71620\n",
        "    100%|██████████| 118/118 [06:17<00:00,  3.20s/it]\n",
        "    100%|██████████| 20/20 [01:04<00:00,  3.25s/it]\n",
        "      0%|          | 0/118 [00:00<?, ?it/s]\n",
        "    Epoch 9: Accuracy 0.97440, Adv Accuracy 0.73010\n",
        "    100%|██████████| 118/118 [06:28<00:00,  3.29s/it]\n",
        "    100%|██████████| 20/20 [01:05<00:00,  3.26s/it]\n",
        "    Epoch 10: Accuracy 0.97580, Adv Accuracy 0.72910"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "39yOpS-62jfk"
      },
      "source": [
        "\n",
        "## 5. (Optional) save the model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3hq7i_m42l6t"
      },
      "source": [
        "os.makedirs(\"models\", exist_ok=True)\n",
        "torch.save(model.state_dict(), f\"models/Net_{num_epochs}_{defense}\")"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}