Adversarial Examples: Fool an Image Classifier, Then Harden It

"""Ferrow Post: the sorting machine reads the handwritten digit on each parcel label with a small neural net.
You find the small pixel changes that make it misread a digit, harden the model against them, and learn to
tell a model that is robust from one whose defence only hides the gradient an attack needs."""
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.datasets import load_digits

torch.set_num_threads(1)


def load_data():
    """8x8 handwritten digits scaled to [0, 1]: 1300 to train on, 497 to test on. Returns tensors."""
    d = load_digits()
    X = (d.data / 16.0).astype(np.float32)
    idx = np.random.default_rng(0).permutation(len(X))
    tr, te = idx[:1300], idx[1300:]
    return (torch.tensor(X[tr]), torch.tensor(d.target[tr]), torch.tensor(X[te]), torch.tensor(d.target[te]))


def new_net():
    """The sorter's classifier: 64 pixels -> 128 hidden units -> 10 digits. Seeded, so it is reproducible."""
    torch.manual_seed(0)
    return nn.Sequential(nn.Linear(64, 128), nn.ReLU(), nn.Linear(128, 10))


def train_standard(X, y, epochs=40):
    """The model in production today, trained on clean digits only."""
    m = new_net()
    opt = torch.optim.Adam(m.parameters(), 1e-2)
    g = torch.Generator().manual_seed(0)
    for _ in range(epochs):
        perm = torch.randperm(len(X), generator=g)
        for i in range(0, len(X), 64):
            b = perm[i:i + 64]
            opt.zero_grad()
            F.cross_entropy(m(X[b]), y[b]).backward()
            opt.step()
    return m


class _Scale(nn.Module):
    def __init__(self, k):
        super().__init__()
        self.k = k

    def forward(self, x):
        return x * self.k


def masked_defence(model):
    """A vendor's 'robustness patch': it multiplies the model's scores by 1000. Predictions do not change."""
    return nn.Sequential(model, _Scale(1000.0))


def accuracy(model, X, y):
    """Share of rows the model classifies correctly."""
    # TODO (Step 1): share of rows where model(X).argmax(1) == y, under torch.no_grad().
    raise NotImplementedError("Step 1: write accuracy()")


def input_gradient(model, X, y, loss_fn=None):
    """Gradient of the loss with respect to the input pixels (default loss: cross-entropy). It points in the
    direction that makes the model most wrong."""
    # TODO (Step 1): gradient of loss_fn (default cross-entropy) w.r.t. a clone of X.
    raise NotImplementedError("Step 1: write input_gradient()")


def fgsm(model, X, y, eps):
    """Fast gradient sign method: move every pixel by eps in the direction of the gradient's sign, then keep
    pixels inside [0, 1]."""
    raise NotImplementedError("fgsm() arrives in Step 2")


def pgd(model, X, y, eps, alpha=None, steps=10, loss_fn=None):
    """Projected gradient descent: `steps` small FGSM steps of size alpha (default eps / 4), each followed by
    projecting back into the eps-box around the original X and then into [0, 1]."""
    raise NotImplementedError("pgd() arrives in Step 3")


def eps_sweep(model, X, y, eps_list, attack):
    """Accuracy under attack(model, X, y, eps) for each eps in eps_list."""
    raise NotImplementedError("eps_sweep() arrives in Step 3")


def adversarial_train(X, y, eps=0.1, epochs=40, steps=5):
    """Train a fresh new_net() on every batch plus its PGD version (eps, `steps` steps), so the model learns
    to classify the attacked digits too."""
    raise NotImplementedError("adversarial_train() arrives in Step 4")


def margin_loss(logits, y):
    """Mean over rows of (best wrong-class score - true-class score). Pushing it up causes a misread, and its
    gradient survives score scaling that flattens cross-entropy."""
    raise NotImplementedError("margin_loss() arrives in Step 5")


def robustness_report(model, X, y, eps=0.1):
    """clean, fgsm, pgd and pgd_margin accuracy at eps; robust = the lowest of the three attacks; sane = the
    cross-entropy PGD attack at eps=0.5 (20 steps) drives accuracy below 0.1, which any working attack must."""
    raise NotImplementedError("robustness_report() arrives in Step 5")