Adversarial Attacks

The attack function gets no true labels; can the frozen classifier provide a usable pseudo-label?

Which objective will reduce confidence in that initial prediction?

How can each gradient step stay inside the per-pixel `L∞` limit?

Project the update back into the `ε` ball and compare the SSIM × accuracy-drop criterion.

Maximize loss for the classifier's initial prediction with three projected sign-gradient steps, staying just inside the 0.3 `L∞` bound.

"""PGD attack for the 2024 Polish AI Olympiad Adversarial Attacks task.

The official notebook supplies ``trained_model.pth`` and calls a one-argument
``perturbe_dataset`` function.  The standalone helper also accepts an explicit
model so it can be measured without changing the notebook.
"""

from pathlib import Path
from typing import Optional

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F


class Net(nn.Module):
    """The exact classifier architecture defined by the organizers."""

    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=0)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=0)
        self.pool = nn.MaxPool2d(kernel_size=3, stride=2)
        self.fc1 = nn.Linear(64 * 11 * 11, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = F.relu(self.conv1(x))
        x = F.relu(self.pool(self.conv2(x)))
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        return F.log_softmax(self.fc2(x), dim=1)


def _load_official_model(
    checkpoint: Optional[str | Path] = None,
    device: Optional[torch.device] = None,
) -> tuple[Net, torch.device]:
    device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
    path = Path(checkpoint) if checkpoint is not None else Path("trained_model.pth")
    model = Net()
    state = torch.load(path, map_location=device, weights_only=True)
    model.load_state_dict(state)
    model.eval().to(device)
    return model, device


def pgd_perturb_dataset(
    original_dataset: np.ndarray,
    model: nn.Module,
    device: Optional[torch.device | str] = None,
    *,
    epsilon: float = 0.299,
    steps: int = 3,
    step_size: float = 0.1,
    batch_size: int = 512,
) -> np.ndarray:
    """Maximize untargeted loss with a projected sign-gradient attack.

    The organizer normalizes inputs independently to [-1, 1].  Output floats
    remain in that same range and are not passed through another normalization.
    ``epsilon`` is kept just below the strict 0.3 pixel-distance limit to avoid
    floating point boundary overshoot in the evaluator.
    """
    images = np.asarray(original_dataset, dtype=np.float32)
    if images.ndim != 3 or images.shape[1:] != (28, 28):
        raise ValueError(f"Expected (N, 28, 28) images, got {images.shape}")
    if not np.isfinite(images).all():
        raise ValueError("Input images must contain only finite values")

    target_device = torch.device(device) if device is not None else next(model.parameters()).device
    was_training = model.training
    model.eval()
    outputs: list[np.ndarray] = []
    with torch.enable_grad():
        for start in range(0, len(images), batch_size):
            clean = torch.from_numpy(images[start : start + batch_size, None]).to(target_device)
            with torch.no_grad():
                # The task interface does not provide ground-truth labels.
                pseudo_labels = model(clean).argmax(dim=1)
            adversarial = clean.detach().clone()
            for _ in range(steps):
                adversarial.requires_grad_(True)
                log_probs = model(adversarial)
                loss = F.nll_loss(log_probs, pseudo_labels)
                gradient = torch.autograd.grad(loss, adversarial, only_inputs=True)[0]
                adversarial = adversarial.detach() + step_size * gradient.sign()
                delta = torch.clamp(adversarial - clean, min=-epsilon, max=epsilon)
                adversarial = torch.clamp(clean + delta, min=-1.0, max=1.0).detach()
            outputs.append(adversarial[:, 0].cpu().numpy())
    if was_training:
        model.train()
    return np.concatenate(outputs, axis=0) if outputs else images.copy()


def perturbe_dataset(
    original_dataset: np.ndarray,
    model: Optional[nn.Module] = None,
    device: Optional[torch.device | str] = None,
    checkpoint: Optional[str | Path] = None,
) -> np.ndarray:
    """Notebook-compatible entry point; loads the official checkpoint if needed."""
    if model is None:
        model, loaded_device = _load_official_model(checkpoint, device)
        device = loaded_device
    return pgd_perturb_dataset(original_dataset, model, device)


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--data", default="data/contest_validation_samples.npy")
    parser.add_argument("--checkpoint", default="trained_model.pth")
    parser.add_argument("--output", default="adversarial_validation.npy")
    args = parser.parse_args()
    raw = np.load(args.data).astype(np.float32) / 255.0
    flat = raw.reshape(len(raw), -1)
    lo = flat.min(axis=1, keepdims=True)
    hi = flat.max(axis=1, keepdims=True)
    normalized = ((2.0 * (flat - lo) / (hi - lo)) - 1.0).reshape(raw.shape)
    model, device = _load_official_model(args.checkpoint)
    adversarial = pgd_perturb_dataset(normalized, model, device)
    np.save(args.output, adversarial)
    print(f"Saved {len(adversarial)} adversarial images to {args.output}")

Download solution.py