{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import numpy as np \nimport pandas as pd","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","collapsed":true,"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":false},"cell_type":"code","source":"import torch\nimport torchvision.datasets as datasets\nimport torchvision.transforms as transforms\nfrom torch.utils.data.sampler import SubsetRandomSampler\nimport torch.utils.data as DataUtils\nimport numpy as np\nimport time\nimport sys\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n# Readymade data loading function\nDATA_ROOT='./MNISTData/'\ndef getMNISTDataLoaders(batchSize=64, nTrain=50000, nVal=10000, nTest=10000):\n  # You can use technically use the same transform instance for all 3 sets\n  assert (60000 - nVal) == nTrain, 'nTrain + nVal must be equal to 60000'\n  trainTransform = transforms.Compose([transforms.ToTensor()])\n  valTransform = transforms.Compose([transforms.ToTensor()])\n  testTransform = transforms.Compose([transforms.ToTensor()])\n  \n  trainSet = datasets.MNIST(root=DATA_ROOT, download=True, train=True, \\\n                           transform=trainTransform)\n  valSet = datasets.MNIST(root=DATA_ROOT, download=True, train=True, \\\n                         transform=valTransform)\n  testSet = datasets.MNIST(root=DATA_ROOT, download=True, train=False, \\\n                                 transform=testTransform)\n  \n  indices = np.arange(0, 60000)\n  np.random.shuffle(indices)\n  \n  trainSampler = SubsetRandomSampler(indices[:nTrain])\n  valSampler = SubsetRandomSampler(indices[nTrain:])\n  testSampler = SubsetRandomSampler(np.arange(0, nTest))\n  \n  trainLoader = DataUtils.DataLoader(trainSet, batch_size=batchSize, \\\n                                   sampler=trainSampler)\n  valLoader = DataUtils.DataLoader(valSet, batch_size=batchSize, \\\n                                  sampler=valSampler)\n  testLoader = DataUtils.DataLoader(testSet, batch_size=batchSize, \\\n                                    sampler=testSampler)\n  return trainLoader, valLoader, testLoader","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class Net(nn.Module):\n    def __init__(self):\n        super(Net, self).__init__()\n        \n        self.conv1 = nn.Sequential(\n            nn.Conv2d(1, 32, 3, padding=1),\n            nn.ReLU(),\n            nn.BatchNorm2d(32),\n            nn.Conv2d(32, 32, 3, stride=2, padding=1),\n            nn.ReLU(),\n            nn.BatchNorm2d(32),\n            nn.MaxPool2d(2, 2),\n            nn.Dropout(0.25)\n        )\n        \n        self.conv2 = nn.Sequential(\n            nn.Conv2d(32, 64, 3, padding=1),\n            nn.ReLU(),\n            nn.BatchNorm2d(64),\n            nn.Conv2d(64, 64, 3, stride=2, padding=1),\n            nn.ReLU(),\n            nn.BatchNorm2d(64),\n            nn.MaxPool2d(2, 2),\n            nn.Dropout(0.25)\n        )\n                \n        self.conv3 = nn.Sequential(\n            nn.Conv2d(64, 128, 3, padding=1),\n            nn.ReLU(),\n            nn.BatchNorm2d(128),\n            nn.MaxPool2d(2, 2),\n            nn.Dropout(0.25)\n        )\n        \n        self.fc = nn.Sequential(\n            nn.Linear(128, 10),\n        )\n                \n        \n    def forward(self, x):\n        x = self.conv1(x)\n        x = self.conv2(x)\n        x = self.conv3(x)\n        \n        x = x.view(x.size(0), -1)\n        return self.fc(x)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"device = 'cuda' if torch.cuda.is_available() else 'cpu'\nprint('Notebook will use PyTorch Device: ' + device.upper())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Utility Progress Bar Function\ndef progress(curr, total, suffix=''):\n  bar_len = 48\n  filled = int(round(bar_len * curr / float(total)))\n  if filled == 0:\n    filled = 1\n  bar = '=' * (filled - 1) + '>' + '-' * (bar_len - filled)\n  sys.stdout.write('\\r[%s] .. %s' % (bar, suffix))\n  sys.stdout.flush()\n  if curr == total:\n    bar = bar_len * '='\n    sys.stdout.write('\\r[%s] .. %s .. Completed\\n' % (bar, suffix))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = Net().to(device)\ncriterion = nn.CrossEntropyLoss()\nlearning_rate = 1e-2\noptimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)\nn_epochs = 1\nlr = 1e-2\nstep = 0\nmodel.train()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_loader, val_loader, test_loader = getMNISTDataLoaders()\nstart_time = time.time()\n\"\"\"\nfor i in range(n_epochs):\n  for j, (images, labels) in enumerate(train_loader):\n    images, labels = images.to(device), labels.to(device)\n    optimizer.zero_grad()\n    logits = model(images)\n    loss = criterion(logits, labels)\n    loss.backward()\n    optimizer.step()\n    if j % 8 == 0:\n      progress(j+1, len(train_loader), 'Batch [{}/{}] Epoch [{}/{}] Loss = {:.3f}'.format(j+1, len(train_loader), i+1, n_epochs, loss.item()))\n    step += 1\nend_time = time.time()\nprint('\\nTotal training steps = {}'.format(step))\nprint('Total time taken = {}'.format(end_time - start_time))\n\"\"\"\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install advertorch > /dev/null\nimport advertorch\nprint(advertorch.__version__)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\n# Evaluating against FGSM attack\nfrom advertorch.attacks import GradientSignAttack\n\"\"\"\n# Documentation for this attack can be found at the link below\n# https://advertorch.readthedocs.io/en/latest/advertorch/attacks.html#advertorch.attacks.GradientSignAttack\nadversary_1 = GradientSignAttack(model, eps=0.3)\ncorrect = 0\nmodel.eval()\nfor j, (images, labels) in enumerate(test_loader):\n  images, labels = images.to(device), labels.to(device)\n  adv_images_1 = adversary_1.perturb(images, labels) # This is extra step as compared to normal clean accuracy testing\n  logits = model(adv_images_1)\n  _, preds = torch.max(logits, 1)\n  correct += (preds == labels).sum().item()\n  progress(j+1, len(test_loader), 'Batch [{}/{}]'.format(j+1, len(test_loader)))\nmodel.train()\nprint('Accuracy on FGSM adversarial samples = {}%'.format(float(correct) * 100 / 10000))\n\"\"\"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\n# Evaluating against IFSGM attack\nfrom advertorch.attacks import LinfBasicIterativeAttack\n\"\"\"\n# Documentation for this attack can be found at the link below\n# https://advertorch.readthedocs.io/en/latest/advertorch/attacks.html#advertorch.attacks.GradientSignAttack\nadversary_2 = LinfBasicIterativeAttack(model, eps=0.1 , nb_iter=40)\ncorrect = 0\nmodel.eval()\nfor j, (images, labels) in enumerate(test_loader):\n  images, labels = images.to(device), labels.to(device)\n  adv_images_2 = adversary_2.perturb(images, labels) # This is extra step as compared to normal clean accuracy testing\n  logits = model(adv_images_2)\n  _, preds = torch.max(logits, 1)\n  correct += (preds == labels).sum().item()\n  progress(j+1, len(test_loader), 'Batch [{}/{}]'.format(j+1, len(test_loader)))\nmodel.train()\nprint('Accuracy on FGSM adversarial samples = {}%'.format(float(correct) * 100 / 10000))\n\"\"\"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\n# Evaluating against PGD attack\nfrom advertorch.attacks import LinfPGDAttack\n\"\"\"\n# Documentation for this attack can be found at the link below\n# https://advertorch.readthedocs.io/en/latest/advertorch/attacks.html#advertorch.attacks.GradientSignAttack\nadversary_3 = LinfPGDAttack(model, eps=0.3 , nb_iter = 40)\ncorrect = 0\nmodel.eval()\nfor j, (images, labels) in enumerate(test_loader):\n  images, labels = images.to(device), labels.to(device)\n  adv_images_3 = adversary_3.perturb(images, labels) # This is extra step as compared to normal clean accuracy testing\n  logits = model(adv_images_3)\n  _, preds = torch.max(logits, 1)\n  correct += (preds == labels).sum().item()\n  progress(j+1, len(test_loader), 'Batch [{}/{}]'.format(j+1, len(test_loader)))\nmodel.train()\nprint('Accuracy on FGSM adversarial samples = {}%'.format(float(correct) * 100 / 10000))\n\"\"\"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"n_epochs = 70\nlr = 1e-2\nstep = 0\nxent_loss = nn.CrossEntropyLoss()\nadv_model = model\nadv_model.train()\noptimizer = torch.optim.SGD(adv_model.parameters(), lr=lr, momentum=0.9)\n\ntrain_loader, val_loader, test_loader = getMNISTDataLoaders()\nstart_time = time.time()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\"\"\"\nAlthough not officially mentioned, making `size_average=False` for the loss \nfunction improves reliability of the result in PyTorch 0.4.0. This is required\nsince we are taking step against the gradient for \"every\" image in the batch.\nSo reducing them to a single value won't cut it.\n\"\"\"\n#training on FSGM\nadvertorch_loss_fn = nn.CrossEntropyLoss(size_average=False)\nfor i in range(n_epochs):\n  for j, (images, labels) in enumerate(train_loader):\n    images, labels = images.to(device), labels.to(device)\n    \"\"\"\n    Creating the adversary :\n    ------------------------\n    Adversarial examples should be typically generated when model parameters are not \n    changing i.e. model parameters are frozen. This step may not be required for very\n    simple linear models, but is a must for models using components such as dropout \n    or batch normalization.\n    \"\"\"\n    adv_model.eval() # Freezes the model parameters\n    \"\"\"\n    The `clip` values here determine the clipping range after taking the adversarial step\n    The clipping is essential to keep the domain of input images within the range\n    MNIST images for this notebook are normalized to [0, 1]. If you're using something else, \n    make sure to modify these values accordingly. The `eps` value decides the magnitude\n    of the attack. For all MNIST models, the threat model advises to stick to maximum eps of 0.3 \n    for input in range [0, 1]\n    \"\"\"\n    fgsm_adversary = GradientSignAttack(adv_model, advertorch_loss_fn, eps=0.3, clip_min=0., \\\n                    clip_max=1., targeted=False)\n    adv_images = fgsm_adversary.perturb(images, labels) # Generate adversarial samples\n    ifgsm_adversary = LinfBasicIterativeAttack(adv_model, advertorch_loss_fn, eps=0.1, clip_min=0., \\\n                    clip_max=1., targeted=False , nb_iter=40)\n    adv_images_2 = ifgsm_adversary.perturb(images, labels) # Generate adversarial samples\n    pgd_adversary = LinfPGDAttack(adv_model, advertorch_loss_fn, eps=0.3, clip_min=0., \\\n                    clip_max=1., targeted=False , nb_iter = 40)\n    adv_images_3 = pgd_adversary.perturb(images, labels)\n    adv_model.train() # Allows model parameters to be changed again\n    optimizer.zero_grad()\n    logits = adv_model(images)\n    loss = criterion(logits, labels)\n    loss.backward()\n    optimizer.step()\n    train_images = adv_images \n    train_labels = labels\n    optimizer.zero_grad()\n    logits = adv_model(train_images)\n    loss = xent_loss(logits, train_labels)\n    loss.backward()\n    optimizer.step()\n    train_images = adv_images_2 \n    train_labels = labels\n    optimizer.zero_grad()\n    logits = adv_model(train_images)\n    loss = xent_loss(logits, train_labels)\n    loss.backward()\n    optimizer.step()\n    train_images = adv_images_3 \n    train_labels = labels\n    optimizer.zero_grad()\n    logits = adv_model(train_images)\n    loss = xent_loss(logits, train_labels)\n    loss.backward()\n    optimizer.step()\n    if j % 8 == 0:\n      progress(j+1, len(train_loader), 'Batch [{}/{}] Epoch [{}/{}] Loss = {:.3f}'.format(j+1, len(train_loader), i+1, n_epochs, loss.item()))\n    step += 1\n\n    end_time = time.time()\nprint('\\nTotal training steps = {}'.format(step))\nprint('Total time taken = {}'.format(end_time - start_time))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\ncorrect = 0\nmodel.eval()\nfor j, (images, labels) in enumerate(test_loader):\n  images, labels = images.to(device), labels.to(device)\n  logits = model(images)\n  _, preds = torch.max(logits, 1)\n  correct += (preds == labels).sum().item()\n  progress(j+1, len(test_loader), 'Batch [{}/{}]'.format(j+1, len(test_loader)))\nmodel.train()\nprint('Accuracy = {}%'.format(float(correct) * 100 / 10000))\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Evaluating against FGSM attack\n\nfrom advertorch.attacks import GradientSignAttack\n# Documentation for this attack can be found at the link below\n# https://advertorch.readthedocs.io/en/latest/advertorch/attacks.html#advertorch.attacks.GradientSignAttack\nadversary_1 = GradientSignAttack(adv_model, eps=0.3)\ncorrect = 0\nadv_model.eval()\nfor j, (images, labels) in enumerate(test_loader):\n  images, labels = images.to(device), labels.to(device)\n  adv_images_1 = adversary_1.perturb(images, labels)\n  logits = adv_model(adv_images_1)\n  _, preds = torch.max(logits, 1)\n  correct += (preds == labels).sum().item()\n  progress(j+1, len(test_loader), 'Batch [{}/{}]'.format(j+1, len(test_loader)))\nadv_model.train()\nprint('Accuracy on FGSM adversarial samples = {}%'.format(float(correct) * 100 / 10000))\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Evaluating against I-FGSM attack\n\nfrom advertorch.attacks import LinfBasicIterativeAttack\n# Documentation for this attack can be found at the link below\n# https://advertorch.readthedocs.io/en/latest/advertorch/attacks.html#advertorch.attacks.GradientSignAttack\nadversary_2 = LinfBasicIterativeAttack(adv_model, eps=0.1 , nb_iter=40)\ncorrect = 0\nadv_model.eval()\nfor j, (images, labels) in enumerate(test_loader):\n  images, labels = images.to(device), labels.to(device)\n  adv_images_2 = adversary_2.perturb(images, labels)\n  logits = adv_model(adv_images_2)\n  _, preds = torch.max(logits, 1)\n  correct += (preds == labels).sum().item()\n  progress(j+1, len(test_loader), 'Batch [{}/{}]'.format(j+1, len(test_loader)))\nadv_model.train()\nprint('Accuracy on FGSM adversarial samples = {}%'.format(float(correct) * 100 / 10000))\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Evaluating against PGD attack\n\nfrom advertorch.attacks import LinfPGDAttack\n# Documentation for this attack can be found at the link below\n# https://advertorch.readthedocs.io/en/latest/advertorch/attacks.html#advertorch.attacks.GradientSignAttack\nadversary_3 = LinfPGDAttack(adv_model, eps=0.3)\ncorrect = 0\nadv_model.eval()\nfor j, (images, labels) in enumerate(test_loader):\n  images, labels = images.to(device), labels.to(device)\n  adv_images_3 = adversary_3.perturb(images, labels)\n  logits = adv_model(adv_images_3)\n  _, preds = torch.max(logits, 1)\n  correct += (preds == labels).sum().item()\n  progress(j+1, len(test_loader), 'Batch [{}/{}]'.format(j+1, len(test_loader)))\nadv_model.train()\nprint('Accuracy on FGSM adversarial samples = {}%'.format(float(correct) * 100 / 10000))\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"torch.save(adv_model.state_dict(), 'model.pt')","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":1}