Adversarial Robustness

"""Adversarial attacks & training — PyTorch + CleverHans.

CleverHans (https://github.com/cleverhans-lab/cleverhans) là thư viện
chuẩn cho adversarial ML. Ở đây ta kết hợp CleverHans cho attack +
PyTorch training loop cho defense.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from cleverhans.torch.attacks.fast_gradient_method import fast_gradient_method
from cleverhans.torch.attacks.projected_gradient_descent import projected_gradient_descent


# ---------------------------------------------------------------------
# 1. Attacks
# ---------------------------------------------------------------------

def fgsm_attack(model: nn.Module, x: torch.Tensor, y: torch.Tensor, eps: float = 0.03):
    """1 bước FGSM — rẻ, hữu ích làm baseline."""
    model.eval()
    # CleverHans API: fast_gradient_method(model, x, eps, norm, ...)
    x_adv = fast_gradient_method(
        model_fn=model,
        x=x,
        eps=eps,
        norm=float("inf"),
        y=y,
        targeted=False,
    )
    return x_adv.detach()


def pgd_attack(model: nn.Module, x: torch.Tensor, y: torch.Tensor,
               eps: float = 0.03, alpha: float = 0.007, steps: int = 10):
    """PGD-K — attack chuẩn để đo robustness."""
    model.eval()
    x_adv = projected_gradient_descent(
        model_fn=model,
        x=x,
        eps=eps,
        eps_iter=alpha,
        nb_iter=steps,
        norm=float("inf"),
        y=y,
        targeted=False,
        rand_init=True,       # khởi tạo ngẫu nhiên trong ball → mạnh hơn
        sanitize=True,
    )
    return x_adv.detach()


# ---------------------------------------------------------------------
# 2. Adversarial training loop (Madry-style)
# ---------------------------------------------------------------------

def train_robust(
    model: nn.Module,
    loader: DataLoader,
    optimizer: torch.optim.Optimizer,
    device: torch.device,
    eps: float = 0.03,
    steps: int = 7,
    clean_weight: float = 0.5,
):
    model.train()
    for x, y in loader:
        x, y = x.to(device), y.to(device)

        # (a) Sinh adversarial batch — mỗi step train đều có attack mới
        x_adv = pgd_attack(model, x, y, eps=eps, alpha=eps/4, steps=steps)
        model.train()  # pgd_attack gọi model.eval()

        # (b) Loss kết hợp clean + adversarial
        logits_clean = model(x)
        logits_adv = model(x_adv)
        loss = (
            clean_weight * F.cross_entropy(logits_clean, y)
            + (1 - clean_weight) * F.cross_entropy(logits_adv, y)
        )

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()


# ---------------------------------------------------------------------
# 3. Đánh giá robustness
# ---------------------------------------------------------------------

@torch.no_grad()
def clean_accuracy(model, loader, device):
    model.eval()
    correct, total = 0, 0
    for x, y in loader:
        x, y = x.to(device), y.to(device)
        correct += (model(x).argmax(1) == y).sum().item()
        total += y.size(0)
    return correct / total


def pgd_accuracy(model, loader, device, eps=0.03, steps=20):
    """Robust accuracy — dùng PGD mạnh (nhiều bước) để đánh giá."""
    model.eval()
    correct, total = 0, 0
    for x, y in loader:
        x, y = x.to(device), y.to(device)
        x_adv = pgd_attack(model, x, y, eps=eps, alpha=eps/10, steps=steps)
        with torch.no_grad():
            correct += (model(x_adv).argmax(1) == y).sum().item()
        total += y.size(0)
    return correct / total


# Thông thường: clean_acc giảm 1-3%, robust_acc tăng từ ~0% → 45-60%.

"""Demo: Tấn công một mô hình CIFAR-10 đã train bằng CleverHans.

Chạy: python demo_cleverhans_attack.py
Kết quả mong đợi: clean_acc ≈ 0.92, fgsm_acc ≈ 0.15, pgd_acc ≈ 0.02
                  (nếu mô hình KHÔNG được adversarial training)
"""
import torch
import torchvision as tv
from cleverhans.torch.attacks.fast_gradient_method import fast_gradient_method
from cleverhans.torch.attacks.projected_gradient_descent import projected_gradient_descent

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 1. Tải CIFAR-10 test set
mean = torch.tensor([0.4914, 0.4822, 0.4465])
std = torch.tensor([0.2470, 0.2435, 0.2616])
transform = tv.transforms.Compose([
    tv.transforms.ToTensor(),
    tv.transforms.Normalize(mean, std),
])
testset = tv.datasets.CIFAR10(root="./data", train=False, download=True, transform=transform)
loader = torch.utils.data.DataLoader(testset, batch_size=128, shuffle=False)

# 2. Tải mô hình — ví dụ ResNet-18 từ torchvision hub
model = torch.hub.load("chenyaofo/pytorch-cifar-models",
                       "cifar10_resnet20", pretrained=True).to(device)
model.eval()

# 3. Hàm đánh giá
def eval_attack(attack_fn, name: str):
    correct, total = 0, 0
    for x, y in loader:
        x, y = x.to(device), y.to(device)
        x_adv = attack_fn(x, y) if attack_fn is not None else x
        with torch.no_grad():
            pred = model(x_adv).argmax(1)
        correct += (pred == y).sum().item()
        total += y.size(0)
    acc = correct / total
    print(f"{name:>10s}  accuracy = {acc:.4f}")
    return acc

# 4. Clean
eval_attack(None, "clean")

# 5. FGSM — 1 bước
eval_attack(
    lambda x, y: fast_gradient_method(model, x, 8/255, float("inf"), y=y),
    "FGSM",
)

# 6. PGD-20 — chuẩn công bố
eval_attack(
    lambda x, y: projected_gradient_descent(
        model, x,
        eps=8/255, eps_iter=2/255, nb_iter=20,
        norm=float("inf"), y=y, rand_init=True,
    ),
    "PGD-20",
)

# 7. AutoAttack (nếu cài autoattack pip package)
# from autoattack import AutoAttack
# aa = AutoAttack(model, norm="Linf", eps=8/255, version="standard")
# x_adv = aa.run_standard_evaluation(x, y)