Machine Learning for CybersecuritySource on GitHub

Unit 7 · Lab 7a

FGSM and PGD against an image classifier

Lab 7a

FGSM and PGD against an image classifier

# On Google Colab, install the scikit-learn version the course was built with.
# On your own computer, install requirements.txt instead; this cell then does nothing.
import subprocess, sys
if "google.colab" in sys.modules:
    subprocess.run([sys.executable, "-m", "pip", "install", "-q", "scikit-learn==1.9.0"], check=True)

Lab 7a. FGSM and PGD against an image classifier

This lab implements two gradient attacks on an image classifier by hand in PyTorch: the fast gradient sign method (FGSM) and projected gradient descent (PGD), both under an L-infinity bound. It trains a small convolutional neural network (CNN) on MNIST, measures accuracy as the attack budget grows, checks that a large enough budget drives accuracy toward zero, and tests whether adversarial examples made against one model fool a second, independently trained model. It runs in a few minutes on CPU.

You will:

  1. train a small CNN on MNIST;
  2. implement FGSM and PGD with random start, step size, projection and clipping;
  3. plot robustness curves, accuracy against the perturbation budget, for both attacks;
  4. confirm the sanity check that accuracy reaches near zero at a large budget;
  5. look at a few adversarial images and their predicted labels;
  6. transfer PGD examples crafted on one model to a second model.

Pixel values are scaled to the interval from 0 to 1, so an L-infinity budget of 0.3 means each pixel may change by at most 0.3 (Lesson 7.2).

import os
import warnings

import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision.datasets import MNIST

warnings.filterwarnings("ignore", category=UserWarning)
torch.set_num_threads(4)
SEED = 20251015
torch.manual_seed(SEED)
rng = np.random.default_rng(SEED)
ROOT = "../../data/downloads"

1. Data and model

We use a balanced subset of MNIST for speed and a compact convolutional network with two convolutional layers, each followed by max pooling, and two fully connected layers.

train = MNIST(root=ROOT, train=True, download=True)
test = MNIST(root=ROOT, train=False, download=True)
Xtr = train.data.numpy().astype(np.float32) / 255.0
ytr = train.targets.numpy()
Xte = test.data.numpy().astype(np.float32) / 255.0
yte = test.targets.numpy()


def balanced(X, y, per_class, seed):
    r = np.random.default_rng(seed)
    idx = np.concatenate([r.choice(np.flatnonzero(y == c), per_class, replace=False) for c in range(10)])
    r.shuffle(idx)
    return torch.from_numpy(X[idx]).unsqueeze(1), torch.from_numpy(y[idx]).long()


Xtr_t, ytr_t = balanced(Xtr, ytr, 500, SEED)
Xte_t, yte_t = balanced(Xte, yte, 100, SEED + 1)
print(f"train {len(Xtr_t)}, test {len(Xte_t)}")


class CNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 16, 3, padding=1)
        self.conv2 = nn.Conv2d(16, 32, 3, padding=1)
        self.fc1 = nn.Linear(32 * 7 * 7, 100)
        self.fc2 = nn.Linear(100, 10)

    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), 2)
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = x.flatten(1)
        return self.fc2(F.relu(self.fc1(x)))


def train_model(seed, epochs=5, lr=2e-3, batch=128):
    torch.manual_seed(seed)
    model = CNN()
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    loss_fn = nn.CrossEntropyLoss()
    r = np.random.default_rng(seed)
    for ep in range(epochs):
        order = r.permutation(len(Xtr_t))
        model.train()
        for i in range(0, len(Xtr_t), batch):
            idx = order[i:i + batch]
            opt.zero_grad()
            loss = loss_fn(model(Xtr_t[idx]), ytr_t[idx])
            loss.backward()
            opt.step()
    model.eval()
    return model


model = train_model(SEED)


@torch.no_grad()
def clean_accuracy(m, X, y):
    return float((m(X).argmax(1) == y).float().mean())


clean_acc = clean_accuracy(model, Xte_t, yte_t)
print(f"clean test accuracy: {clean_acc:.3f}")
train 5000, test 1000
clean test accuracy: 0.948

2. FGSM and PGD by hand

FGSM takes one step in the direction that increases the loss, scaled by the budget, then clips to the valid pixel range,

x_adv = clip(x + eps * sign(grad_x L(x, y)), 0, 1).

Here x is the input image with pixels in the interval from 0 to 1, y is its true label, L is the cross-entropy loss of the model, grad_x L is the gradient of that loss with respect to the pixels, sign is taken pixel by pixel, and eps is the L-infinity budget, as in Lesson 7.2.

PGD starts at a random point inside the L-infinity ball, then repeats a small gradient-sign step and projects the result back into the ball intersected with the valid range. Both use the gradient of the loss with respect to the input, not the parameters.

loss_fn = nn.CrossEntropyLoss()


def fgsm(model, X, y, eps):
    X = X.clone().detach().requires_grad_(True)
    loss = loss_fn(model(X), y)
    grad, = torch.autograd.grad(loss, X)
    X_adv = X + eps * grad.sign()
    return X_adv.clamp(0, 1).detach()


def pgd_run(model, X, y, eps, alpha, steps):
    # random start, uniform in the L-infinity ball, kept inside the valid range
    delta = (torch.rand_like(X) * 2 - 1) * eps
    Xa = (X + delta).clamp(0, 1).detach()
    for _ in range(steps):
        Xa.requires_grad_(True)
        loss = loss_fn(model(Xa), y)
        grad, = torch.autograd.grad(loss, Xa)
        Xa = Xa.detach() + alpha * grad.sign()
        # projection onto the intersection of the eps-ball and [0, 1]
        Xa = torch.min(torch.max(Xa, X - eps), X + eps).clamp(0, 1).detach()
    return Xa


def pgd(model, X, y, eps, alpha=None, steps=20, restarts=3):
    """PGD with random restarts. An input counts as broken if any restart misclassifies it;
    later restarts are run only on the inputs that are still classified correctly."""
    if alpha is None:
        alpha = 2.5 * eps / steps
    best = pgd_run(model, X, y, eps, alpha, steps)
    with torch.no_grad():
        wrong = model(best).argmax(1) != y
    for _ in range(max(1, restarts) - 1):
        remaining = (~wrong).nonzero(as_tuple=True)[0]
        if len(remaining) == 0:
            break
        Xa = pgd_run(model, X[remaining], y[remaining], eps, alpha, steps)
        with torch.no_grad():
            new_wrong = model(Xa).argmax(1) != y[remaining]
        best[remaining] = Xa
        wrong[remaining] = new_wrong
    return best

3. Robustness curves

For a range of budgets we craft adversarial examples on the test set and measure accuracy. PGD uses 10 steps of size alpha = 2.5 eps / 10 and 3 random restarts, and an input counts as broken if any restart misclassifies it. The attacks run in mini-batches to bound memory.

EPS = [0.0, 0.05, 0.1, 0.15, 0.2, 0.3, 0.4]


def attacked_accuracy(model, attack, eps, batch=200, **kw):
    if eps == 0.0:
        return clean_accuracy(model, Xte_t, yte_t)
    correct = 0
    for i in range(0, len(Xte_t), batch):
        Xb, yb = Xte_t[i:i + batch], yte_t[i:i + batch]
        Xadv = attack(model, Xb, yb, eps, **kw)
        with torch.no_grad():
            correct += int((model(Xadv).argmax(1) == yb).sum())
    return correct / len(Xte_t)


fgsm_curve = [attacked_accuracy(model, fgsm, e) for e in EPS]
pgd_curve = [attacked_accuracy(model, pgd, e, steps=10, restarts=3) for e in EPS]
for e, af, ap in zip(EPS, fgsm_curve, pgd_curve):
    print(f"eps={e:.2f}  FGSM acc {af:.3f}   PGD acc {ap:.3f}")
eps=0.00  FGSM acc 0.948   PGD acc 0.948
eps=0.05  FGSM acc 0.869   PGD acc 0.855
eps=0.10  FGSM acc 0.725   PGD acc 0.619
eps=0.15  FGSM acc 0.478   PGD acc 0.239
eps=0.20  FGSM acc 0.224   PGD acc 0.019
eps=0.30  FGSM acc 0.010   PGD acc 0.000
eps=0.40  FGSM acc 0.000   PGD acc 0.000
fig, ax = plt.subplots(figsize=(6, 3))
ax.plot(EPS, fgsm_curve, "s--", label="FGSM (one step)")
ax.plot(EPS, pgd_curve, "o-", label="PGD (10 steps, 3 restarts)")
ax.set_xlabel("budget eps (L-infinity, pixels from 0 to 1)")
ax.set_ylabel("test accuracy")
ax.legend()
fig.tight_layout()
plt.show()
Plot produced by the cell above

PGD is the stronger attack at the same budget, and in the table and plot above its accuracy is at or below FGSM's at every budget. The sanity check of Lesson 7.3 is the accuracy at the largest budget. Within the valid pixel range, a strong attack should drive accuracy close to zero.

pgd_at_03 = pgd_curve[EPS.index(0.3)]
pgd_at_04 = pgd_curve[EPS.index(0.4)]
fgsm_at_03 = fgsm_curve[EPS.index(0.3)]
print(f"PGD accuracy at eps=0.3: {pgd_at_03:.3f}")
print(f"PGD accuracy at eps=0.4: {pgd_at_04:.3f}  (sanity check: near zero)")
PGD accuracy at eps=0.3: 0.000
PGD accuracy at eps=0.4: 0.000  (sanity check: near zero)

4. Example adversarial images

We craft PGD examples at eps = 0.2 on a few correctly classified digits, show the clean and perturbed images with their predicted labels, and save them for the lesson figure.

with torch.no_grad():
    ok = (model(Xte_t).argmax(1) == yte_t)
ex_idx = ok.nonzero(as_tuple=True)[0][:6]
ex_clean = Xte_t[ex_idx]
ex_labels = yte_t[ex_idx]
ex_adv = pgd(model, ex_clean, ex_labels, eps=0.2, steps=20, restarts=1)
with torch.no_grad():
    clean_pred = model(ex_clean).argmax(1).numpy()
    adv_pred = model(ex_adv).argmax(1).numpy()
print("clean predictions:", clean_pred.tolist())
print("adversarial predictions:", adv_pred.tolist())
clean predictions:
 [5, 1, 3, 5, 4, 4]
adversarial predictions: [3, 8, 5, 9, 9, 8]
fig, axes = plt.subplots(2, len(ex_idx), figsize=(9, 3.2))
for j in range(len(ex_idx)):
    axes[0, j].imshow(ex_clean[j, 0], cmap="gray_r", vmin=0, vmax=1)
    axes[0, j].set_title(f"clean, pred {clean_pred[j]}", fontsize=8)
    axes[1, j].imshow(ex_adv[j, 0], cmap="gray_r", vmin=0, vmax=1)
    axes[1, j].set_title(f"PGD, pred {adv_pred[j]}", fontsize=8)
for a in axes.ravel():
    a.axis("off")
plt.show()
Plot produced by the cell above

5. Transfer to a second model

We train a second model with a different seed and feed it the PGD examples crafted on the first model, without using the second model's gradients. The drop in its accuracy measures how well the attack transfers between independently trained models.

model2 = train_model(SEED + 100)
clean_acc2 = clean_accuracy(model2, Xte_t, yte_t)

TRANSFER_EPS = 0.2
correct_src = correct_dst = 0
for i in range(0, len(Xte_t), 200):
    Xb, yb = Xte_t[i:i + 200], yte_t[i:i + 200]
    Xadv = pgd(model, Xb, yb, TRANSFER_EPS, steps=20, restarts=2)
    with torch.no_grad():
        correct_src += int((model(Xadv).argmax(1) == yb).sum())
        correct_dst += int((model2(Xadv).argmax(1) == yb).sum())
src_acc = correct_src / len(Xte_t)
transfer_acc = correct_dst / len(Xte_t)
print(f"model 2 clean accuracy: {clean_acc2:.3f}")
print(f"on PGD(eps={TRANSFER_EPS}) examples crafted for model 1: model 1 {src_acc:.3f}, model 2 {transfer_acc:.3f}")
model 2 clean accuracy: 0.956
on PGD(eps=0.2) examples crafted for model 1: model 1 0.018, model 2 0.317

The examples that defeat the first model lower the second model's accuracy too, though less than the first model's, which is the transferability of Lesson 7.1. The attack is white-box against the first model. Crafting the examples uses neither the gradients nor the outputs of the second model, but the first model shares the second model's architecture and training data, so in the terms of Lesson 7.1 the attacker has gray-box knowledge of the second model (everything except its parameters).

Try this

Results summary

These are the numbers quoted in the lesson text.

results = {
    "n_train": int(len(Xtr_t)),
    "n_test": int(len(Xte_t)),
    "clean_acc": float(clean_acc),
    "fgsm_acc_01": float(fgsm_curve[EPS.index(0.1)]),
    "fgsm_acc_03": float(fgsm_at_03),
    "pgd_acc_01": float(pgd_curve[EPS.index(0.1)]),
    "pgd_acc_03": float(pgd_at_03),
    "pgd_acc_04": float(pgd_at_04),
    "clean_acc2": float(clean_acc2),
    "transfer_src_acc": float(src_acc),
    "transfer_dst_acc": float(transfer_acc),
    "transfer_eps": TRANSFER_EPS,
}
for k, v in results.items():
    print(f"{k:<18} {v}")

export = {
    "eps": EPS,
    "fgsm_curve": [float(x) for x in fgsm_curve],
    "pgd_curve": [float(x) for x in pgd_curve],
    "example_clean": ex_clean.squeeze(1).numpy().tolist(),
    "example_adv": ex_adv.squeeze(1).numpy().tolist(),
    "clean_pred": clean_pred.tolist(),
    "adv_pred": adv_pred.tolist(),
    "true_labels": ex_labels.numpy().tolist(),
}

if os.environ.get("COURSE_RESULTS"):
    import json
    with open(os.environ["COURSE_RESULTS"], "w") as f:
        json.dump({"results": results, "data": export}, f)
n_train            5000
n_test             1000
clean_acc          0.9480000138282776
fgsm_acc_01        0.725
fgsm_acc_03        0.01
pgd_acc_01         0.619
pgd_acc_03         0.0
pgd_acc_04         0.0
clean_acc2         0.9559999704360962
transfer_src_acc   0.018
transfer_dst_acc   0.317
transfer_eps       0.2