# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python [conda env:root] *
#     language: python
#     name: conda-root-py
# ---

# + [markdown] id="view-in-github" colab_type="text"
# <a href="https://colab.research.google.com/github/spKay18/interactive_latencies/blob/master/part_1_training.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>

# + id="Z4ASX8AZ9uCy" colab_type="code" outputId="ef71ba4a-1bb8-4dbc-b6c1-f97c27cba03c" colab={"base_uri": "https://localhost:8080/", "height": 121}
from google.colab import drive
drive.mount('/content/drive')

# + id="fzU3UIB_9qge" colab_type="code" outputId="0a9526d1-569d-4a1e-ad29-52e646c02bdb" colab={"base_uri": "https://localhost:8080/", "height": 50}
# %load_ext autoreload
# %autoreload 2
import torch
from torch.utils.data import DataLoader
from torch.optim import Adam
from matplotlib import pyplot as plt
from utils import get_mnist_data
from models import ConvNN
from training_and_evaluation import train_model, predict_model
from attacks import fast_gradient_attack
from torch.nn.functional import cross_entropy
import os
if not os.path.isdir("models"):
    os.mkdir("models")

# + [markdown] id="vHg3twW39qgh" colab_type="text"
# # Project 2, part 1: Creating adversarial examples (50 pt)
# In this notebook we train a basic convolutional neural network on MNIST and craft adversarial examples via gradient descent.
#
# ## Your task
# Complete the missing code in the respective files, i.e. `training_and_evaluation.py`, `attacks.py`, and this notebook. Make sure that all the functions follow the provided specification, i.e. the output of the function exactly matches the description in the docstring. 
#
# Specifically, for this part you will have to implement the following functions / classes:  
# **`training_and_evaluation.py`**:
# * `train_model` (15pt)
# * `predict_model` (10pt)
#
# **`attacks.py`**:
# * `fast_gradient_attack` (15pt)
#
# **This notebook**
# * Cells in the Qualitative Evaluation section. (10pt)
#
# ## General remarks
# Do not add or modify any code outside of the following comment blocks, or where otherwise explicitly stated.
#
# ``` python
# ##########################################################
# # YOUR CODE HERE
# ...
# ##########################################################
# ```
# After you fill in all the missing code, restart the kernel and re-run all the cells in the notebook.
#
# The following things are **NOT** allowed:
# - Using additional `import` statements
# - Copying / reusing code from other sources (e.g. code by other students)
#
# If you plagiarise even for a single project task, you won't be eligible for the bonus this semester.

# + id="3QBG_8BY9qgh" colab_type="code" colab={}
mnist_trainset = get_mnist_data(train=True)
mnist_testset = get_mnist_data(train=False)

use_cuda = torch.cuda.is_available() #and False

model = ConvNN()
if use_cuda:
    model = model.cuda()

epochs = 1
batch_size = 128
test_batch_size = 1000  # feel free to change this
lr = 1e-3

opt = Adam(model.parameters(), lr=lr)


# + id="Z9x8iL719qgm" colab_type="code" colab={}
def loss_function(x, y, model):
    logits = model(x).cpu()
    loss = cross_entropy(logits, y)
    return loss, logits


# + [markdown] id="Jedehy5u9qgo" colab_type="text"
# Implement the `train_model` function in the file `training_and_evaluation.py`.

# + id="EyCqUtZp9qgp" colab_type="code" outputId="0ce4c823-274c-4f23-e233-f14f4f870269" colab={"base_uri": "https://localhost:8080/", "height": 66, "referenced_widgets": ["65fba63a40e0450b877e6f45f9faec58", "ef6d7c6be40040059706e45506fd51c2", "daf47e1ac07543248f1ca750002ca7a1", "042430a5a53944aa893f11ac0eaf585e", "cd9c4dbe73fd44189a42c9327b51e1f8", "06c2049d99534deca780488d234ee162", "48a4bb8f63c04c68b02e1c61a1dc02ab", "397400e9880e4a0b94c20d4504a8ceb9"]}
losses, accuracies = train_model(model, mnist_trainset, batch_size=batch_size, loss_function=loss_function, optimizer=opt)

# + id="eN7JyxjI9qgs" colab_type="code" colab={}
torch.save(model.state_dict(), "models/standard_training.checkpoint")

# + id="oQnm43JU9qgu" colab_type="code" outputId="cf119e63-4466-40b2-90c7-4ce85b57c800" colab={"base_uri": "https://localhost:8080/", "height": 34}
model.load_state_dict(torch.load("models/standard_training.checkpoint", map_location="cpu"))

# + id="FRyIj6Y49qgw" colab_type="code" outputId="3bbb2684-72d7-4dd8-8017-fc733676911f" colab={"base_uri": "https://localhost:8080/", "height": 225}
fig = plt.figure(figsize=(10,3))
plt.subplot(121)
plt.plot(losses)
plt.xlabel("Iteration")
plt.ylabel("Training Loss")
plt.subplot(122)
plt.plot(accuracies)
plt.xlabel("Iteration")
plt.ylabel("Training Accuracy")
plt.show()

# + [markdown] id="neT5Ai_w9qgz" colab_type="text"
# Implement the `predict_model` function in the file `training_and_evaluation.py`.

# + id="QQRgMdCL9qg1" colab_type="code" outputId="be1366dd-0876-4ea4-a265-87e2496ec1cb" colab={"base_uri": "https://localhost:8080/", "height": 82, "referenced_widgets": ["cbd305b47c934374ad1eee0bbb826eba", "e68dbf61eae843aea1207f363aca6cc1", "03acebed2090403dadd354f48d68823d", "c65e32aa0980438d9917cb4a37f4aad9", "f87adb0b11314662b699a91d8485afd2", "7432175e74c341af87b8c200320579e6", "c6105b926251498ba30893b9e128818a", "<KEY>"]}
clean_accuracy = predict_model(model, mnist_testset, batch_size=test_batch_size, attack_function=None)

# + [markdown] id="VXjvkkpC9qg4" colab_type="text"
# ### Creating adversarial examples
# #### $L_2$-bounded attacks
# Fist, craft adversarial perturbations that have a $L_2$ norm of $ \| \tilde{\mathbf{x}} - \mathbf{x} \|_2 = \epsilon$ with $\epsilon=5$.
#
# #### $L_\infty$-bounded attacks
# Afterwards, craft adversarial perturbations with $L_\infty$ norm of $ \| \tilde{\mathbf{x}} - \mathbf{x} \|_\infty = \epsilon$ with $\epsilon=0.3$.
#
# For this you need to implement `predict_model` in the file `training_and_evaluation.py` and `fast_gradient_attack` in `attacks.py`. See the docstring comments there.

# + id="Ix8KWG3A9qg5" colab_type="code" colab={}
attack_args_l2 = {"epsilon": 5, "norm": "2"}
attack_args_linf = {"epsilon": 0.3, "norm": "inf"}

# + [markdown] id="v6wlHdA39qg7" colab_type="text"
# ### Qualitative evaluation
#
# First, craft adversarial examples for 10 randomly selected test samples and inspect them by plotting them.

# + [markdown] id="gd9p_pS-9qg8" colab_type="text"
# $L_2$ attack:

# + id="Tolo_zEf9qg8" colab_type="code" colab={}
test_loader = DataLoader(mnist_testset, batch_size = 10, shuffle=True)
x,y = next(iter(test_loader))
##########################################################
# YOUR CODE HERE
...
##########################################################

# + [markdown] id="TK6kdjtv9qg-" colab_type="text"
# $L_\infty$ attack:

# + id="Noq_fhdY9qg_" colab_type="code" colab={}
##########################################################
# YOUR CODE HERE
...
##########################################################

# + [markdown] id="H3bK109D9qhB" colab_type="text"
# Visualize the adversarial examples and the model's prediction on them:

# + id="zhzM5blN9qhB" colab_type="code" outputId="9f4555d9-2d07-474d-903b-83844e60af64" colab={}
for ix in range(len(x)):
    plt.subplot(131)
    plt.imshow(x[ix,0].detach().cpu(), cmap="gray")
    plt.title(f"Label: {y[ix]}")
    plt.subplot(132)
    plt.imshow(x_pert_l2[ix,0].detach().cpu(), cmap="gray")
    plt.title(f"Predicted: {y_pert_l2[ix]}")
    
    plt.subplot(133)
    plt.imshow(x_pert_linf[ix,0].detach().cpu(), cmap="gray")
    plt.title(f"Predicted: {y_pert_linf[ix]}")
    plt.show()

# + [markdown] id="u6435w719qhE" colab_type="text"
# ### Quantitative evaluation
# Perturb each test sample and compare the clean and perturbed accuracies.

# + [markdown] id="dY4SJcI-9qhE" colab_type="text"
# $L_2$ perturbations:

# + id="tBOvScuJ9qhF" colab_type="code" outputId="e640bb71-a633-4e24-e67d-51dcccb1693a" colab={"referenced_widgets": ["3b77f76347a142aeba5b353232d67865"]}
perturbed_accuracy_l2 = predict_model(model, mnist_testset, batch_size=test_batch_size, attack_function=fast_gradient_attack, attack_args=attack_args_l2)

# + [markdown] id="ugzjV0xQ9qhK" colab_type="text"
# $L_\infty$ perturbations:

# + id="3ulxn0eq9qhK" colab_type="code" outputId="3f3c5ae1-e19e-4982-e24e-53912e0d7671" colab={"referenced_widgets": ["3e66590ea87741229bed318d06af52da"]}
perturbed_accuracy_linf = predict_model(model, mnist_testset, batch_size=test_batch_size, attack_function=fast_gradient_attack, attack_args=attack_args_linf)

# + [markdown] id="lGDTnke79qhM" colab_type="text"
# Your values for `clean_accuracy` and `perturbed_accuracy` should roughly match the ones below, even though they will of course not be identical.

# + id="RrVWtOg99qhM" colab_type="code" outputId="344e57c0-c75c-46ad-ab7a-0786e53ae658" colab={}
clean_accuracy

# + id="PuOXWGBm9qhP" colab_type="code" outputId="66feec1e-c0ed-4f0b-abdd-7b95c7590736" colab={}
perturbed_accuracy_l2

# + id="UOPT8V6q9qhS" colab_type="code" outputId="1153cef7-168f-4d1a-8132-1fb33b929dca" colab={}
perturbed_accuracy_linf

# + [markdown] id="hP-zt9FU9qhU" colab_type="text"
# #### In the remaining parts of this project we will be focusing on **$L_2$-based attacks only**.

StarcoderdataJupyterScriptsDedupFiltered