Jeremy Janella
582 lines
19 KiB
Python
582 lines
19 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Adversarial Robustness Evaluator
|
|
|
|
This standalone evaluator tests your adversarially-trained model against
|
|
pre-generated FGSM and I-FGSM attacks. Run this after training to measure
|
|
your model's robustness metrics.
|
|
|
|
Usage:
|
|
python evaluate_robustness.py --model-path robust_model.safetensors
|
|
python evaluate_robustness.py --model-path models/ --compare
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Dict, List, Tuple, Optional
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Configuration
|
|
# ---------------------------------------------------------------------------
|
|
|
|
SCRIPT_DIR = Path(__file__).resolve().parent
|
|
|
|
# Primary epsilon for main results
|
|
PRIMARY_EPSILON = 0.3
|
|
|
|
# Files to skip when scanning directories for models
|
|
SKIP_FILES = {"baseline_model.safetensors", "adv_examples.safetensors"}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Model Definition (lazy import)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def create_lenet5():
|
|
"""Create a LeNet-5 model instance."""
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
class LeNet5(nn.Module):
|
|
"""Classic LeNet-5 architecture for MNIST classification."""
|
|
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.conv1 = nn.Conv2d(1, 6, kernel_size=5, padding=2)
|
|
self.conv2 = nn.Conv2d(6, 16, kernel_size=5)
|
|
self.fc1 = nn.Linear(16 * 5 * 5, 120)
|
|
self.fc2 = nn.Linear(120, 84)
|
|
self.fc3 = nn.Linear(84, 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.view(-1, 16 * 5 * 5)
|
|
x = F.relu(self.fc1(x))
|
|
x = F.relu(self.fc2(x))
|
|
x = self.fc3(x)
|
|
return x
|
|
|
|
return LeNet5()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Safetensors Helper Functions
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def load_adversarial_examples_safetensors(filepath: Path) -> Dict:
|
|
"""Load adversarial examples from safetensors format."""
|
|
from safetensors import safe_open
|
|
from safetensors.torch import load_file
|
|
|
|
tensors = load_file(str(filepath))
|
|
|
|
# Extract metadata
|
|
with safe_open(str(filepath), framework="pt") as f:
|
|
metadata = f.metadata()
|
|
|
|
epsilon = float(metadata["epsilon"])
|
|
epsilon_spread = json.loads(metadata["epsilon_spread"])
|
|
|
|
# Reconstruct epsilon-keyed dicts
|
|
fgsm_by_epsilon = {}
|
|
ifgsm_by_epsilon = {}
|
|
for key, tensor in tensors.items():
|
|
if key.startswith("fgsm_eps_"):
|
|
eps = float(key.replace("fgsm_eps_", ""))
|
|
fgsm_by_epsilon[eps] = tensor
|
|
elif key.startswith("ifgsm_eps_"):
|
|
eps = float(key.replace("ifgsm_eps_", ""))
|
|
ifgsm_by_epsilon[eps] = tensor
|
|
|
|
return {
|
|
"clean_images": tensors["clean_images"],
|
|
"clean_labels": tensors["clean_labels"],
|
|
"fgsm_images": fgsm_by_epsilon[epsilon],
|
|
"ifgsm_images": ifgsm_by_epsilon[epsilon],
|
|
"epsilon": epsilon,
|
|
"epsilon_spread": epsilon_spread,
|
|
"fgsm_by_epsilon": fgsm_by_epsilon,
|
|
"ifgsm_by_epsilon": ifgsm_by_epsilon,
|
|
}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# File Discovery
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def find_models(path: Path) -> List[Path]:
|
|
"""Find model files to evaluate."""
|
|
if path.is_file() and path.suffix == ".safetensors":
|
|
return [path]
|
|
|
|
if path.is_dir():
|
|
models = []
|
|
for f in sorted(path.glob("*.safetensors")):
|
|
if f.name not in SKIP_FILES:
|
|
models.append(f)
|
|
return models
|
|
|
|
return []
|
|
|
|
|
|
def find_baseline(model_path: Path) -> Optional[Path]:
|
|
"""Find baseline_model.safetensors relative to the model path."""
|
|
if model_path.is_file():
|
|
search_dir = model_path.parent
|
|
else:
|
|
search_dir = model_path
|
|
|
|
for check_dir in [search_dir, search_dir.parent, SCRIPT_DIR]:
|
|
baseline = check_dir / "baseline_model.safetensors"
|
|
if baseline.exists():
|
|
return baseline
|
|
|
|
return None
|
|
|
|
|
|
def find_adv_examples(model_path: Path) -> Optional[Path]:
|
|
"""Find adv_examples.safetensors relative to the model path."""
|
|
if model_path.is_file():
|
|
search_dir = model_path.parent
|
|
else:
|
|
search_dir = model_path
|
|
|
|
for check_dir in [search_dir, search_dir.parent, SCRIPT_DIR]:
|
|
adv = check_dir / "adv_examples.safetensors"
|
|
if adv.exists():
|
|
return adv
|
|
|
|
return None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Evaluation Functions
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def load_model(model_path: Path):
|
|
"""Load a trained model from safetensors checkpoint."""
|
|
from safetensors.torch import load_file
|
|
|
|
model = create_lenet5()
|
|
state_dict = load_file(str(model_path))
|
|
model.load_state_dict(state_dict)
|
|
model.eval()
|
|
return model
|
|
|
|
|
|
def load_adversarial_examples(adv_path: Path) -> Dict:
|
|
"""Load pre-generated adversarial examples from safetensors."""
|
|
return load_adversarial_examples_safetensors(adv_path)
|
|
|
|
|
|
def evaluate_accuracy(model, images, labels, device) -> Tuple[float, int, int]:
|
|
"""Compute accuracy and return (accuracy%, correct, total)."""
|
|
import torch
|
|
|
|
model.eval()
|
|
with torch.no_grad():
|
|
images = images.to(device)
|
|
labels = labels.to(device)
|
|
outputs = model(images)
|
|
_, predicted = outputs.max(1)
|
|
correct = predicted.eq(labels).sum().item()
|
|
total = labels.size(0)
|
|
|
|
return 100.0 * correct / total, correct, total
|
|
|
|
|
|
def get_misclassified_samples(
|
|
model, images, labels, device, max_samples: int = 5
|
|
) -> list:
|
|
"""Get indices and predictions of misclassified samples."""
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
model.eval()
|
|
with torch.no_grad():
|
|
images = images.to(device)
|
|
labels = labels.to(device)
|
|
outputs = model(images)
|
|
_, predicted = outputs.max(1)
|
|
wrong = predicted != labels
|
|
wrong_indices = wrong.nonzero(as_tuple=True)[0][:max_samples]
|
|
|
|
failures = []
|
|
for idx in wrong_indices:
|
|
failures.append(
|
|
{
|
|
"index": idx.item(),
|
|
"true_label": labels[idx].item(),
|
|
"predicted": predicted[idx].item(),
|
|
"confidence": F.softmax(outputs[idx], dim=0).max().item(),
|
|
}
|
|
)
|
|
|
|
return failures
|
|
|
|
|
|
def evaluate_model_full(model_path: Path, adv_data: Dict, device) -> Dict:
|
|
"""Evaluate a model across all epsilon values and return full results."""
|
|
clean_images = adv_data["clean_images"]
|
|
clean_labels = adv_data["clean_labels"]
|
|
epsilon = adv_data["epsilon"]
|
|
|
|
# Check for epsilon spread vs single epsilon
|
|
has_spread = "epsilon_spread" in adv_data and "fgsm_by_epsilon" in adv_data
|
|
if has_spread:
|
|
epsilon_spread = adv_data["epsilon_spread"]
|
|
fgsm_by_epsilon = adv_data["fgsm_by_epsilon"]
|
|
ifgsm_by_epsilon = adv_data["ifgsm_by_epsilon"]
|
|
else:
|
|
epsilon_spread = [epsilon]
|
|
fgsm_by_epsilon = {epsilon: adv_data["fgsm_images"]}
|
|
ifgsm_by_epsilon = {epsilon: adv_data["ifgsm_images"]}
|
|
|
|
model = load_model(model_path)
|
|
model.to(device)
|
|
|
|
# Evaluate clean accuracy
|
|
clean_acc, clean_correct, clean_total = evaluate_accuracy(
|
|
model, clean_images, clean_labels, device
|
|
)
|
|
|
|
# Evaluate across all epsilon values
|
|
fgsm_results = {}
|
|
ifgsm_results = {}
|
|
for eps in epsilon_spread:
|
|
fgsm_acc, _, _ = evaluate_accuracy(
|
|
model, fgsm_by_epsilon[eps], clean_labels, device
|
|
)
|
|
ifgsm_acc, _, _ = evaluate_accuracy(
|
|
model, ifgsm_by_epsilon[eps], clean_labels, device
|
|
)
|
|
fgsm_results[eps] = fgsm_acc
|
|
ifgsm_results[eps] = ifgsm_acc
|
|
|
|
return {
|
|
"name": model_path.name,
|
|
"clean_acc": clean_acc,
|
|
"clean_correct": clean_correct,
|
|
"clean_total": clean_total,
|
|
"fgsm_results": fgsm_results,
|
|
"ifgsm_results": ifgsm_results,
|
|
"epsilon_spread": epsilon_spread,
|
|
"model": model,
|
|
}
|
|
|
|
|
|
def print_model_results(results: Dict, baseline: Optional[Dict] = None) -> None:
|
|
"""Print detailed results for a single model."""
|
|
name = results["name"]
|
|
clean_acc = results["clean_acc"]
|
|
fgsm_results = results["fgsm_results"]
|
|
ifgsm_results = results["ifgsm_results"]
|
|
epsilon_spread = results["epsilon_spread"]
|
|
|
|
print(f"\n{'=' * 70}")
|
|
print(f"Model: {name}")
|
|
print(f"{'=' * 70}")
|
|
|
|
# Clean accuracy
|
|
print(
|
|
f"\nClean Accuracy: {clean_acc:.1f}% ({results['clean_correct']}/{results['clean_total']})"
|
|
)
|
|
|
|
# Epsilon spread table
|
|
print(
|
|
f"\n{'Epsilon':<10} {'FGSM':>10} {'I-FGSM':>10} {'FGSM Attack':>12} {'I-FGSM Attack':>14}"
|
|
)
|
|
print("-" * 60)
|
|
|
|
for eps in epsilon_spread:
|
|
fgsm_acc = fgsm_results[eps]
|
|
ifgsm_acc = ifgsm_results[eps]
|
|
fgsm_attack = 100.0 - fgsm_acc
|
|
ifgsm_attack = 100.0 - ifgsm_acc
|
|
marker = " *" if eps == PRIMARY_EPSILON else ""
|
|
print(
|
|
f"{eps:<10.2f} {fgsm_acc:>9.1f}% {ifgsm_acc:>9.1f}% {fgsm_attack:>11.1f}% {ifgsm_attack:>13.1f}%{marker}"
|
|
)
|
|
|
|
print("-" * 60)
|
|
print("* Primary epsilon for summary metrics")
|
|
|
|
# Summary at primary epsilon
|
|
primary_fgsm = fgsm_results[PRIMARY_EPSILON]
|
|
primary_ifgsm = ifgsm_results[PRIMARY_EPSILON]
|
|
|
|
print(f"\nSummary (ε={PRIMARY_EPSILON}):")
|
|
print(f" Clean: {clean_acc:5.1f}%")
|
|
print(
|
|
f" FGSM: {primary_fgsm:5.1f}% (attack success: {100 - primary_fgsm:5.1f}%)"
|
|
)
|
|
print(
|
|
f" I-FGSM: {primary_ifgsm:5.1f}% (attack success: {100 - primary_ifgsm:5.1f}%)"
|
|
)
|
|
|
|
# Comparison with baseline if provided
|
|
if baseline:
|
|
base_fgsm = baseline["fgsm_results"][PRIMARY_EPSILON]
|
|
base_ifgsm = baseline["ifgsm_results"][PRIMARY_EPSILON]
|
|
|
|
fgsm_improve = primary_fgsm - base_fgsm
|
|
ifgsm_improve = primary_ifgsm - base_ifgsm
|
|
|
|
print(f"\nImprovement over baseline:")
|
|
print(
|
|
f" FGSM: {'+' if fgsm_improve >= 0 else ''}{fgsm_improve:.1f}% ({base_fgsm:.1f}% → {primary_fgsm:.1f}%)"
|
|
)
|
|
print(
|
|
f" I-FGSM: {'+' if ifgsm_improve >= 0 else ''}{ifgsm_improve:.1f}% ({base_ifgsm:.1f}% → {primary_ifgsm:.1f}%)"
|
|
)
|
|
|
|
|
|
def print_comparison_summary(
|
|
all_results: List[Dict], baseline: Optional[Dict] = None
|
|
) -> None:
|
|
"""Print a full epsilon spread comparison table of all models."""
|
|
epsilon_spread = all_results[0]["epsilon_spread"]
|
|
|
|
print(f"\n{'=' * 90}")
|
|
print("COMPARISON SUMMARY - Full Epsilon Spread")
|
|
print(f"{'=' * 90}")
|
|
|
|
# Build model list (baseline first if present)
|
|
models_to_show = []
|
|
if baseline:
|
|
models_to_show.append(("baseline", baseline))
|
|
for r in all_results:
|
|
models_to_show.append((r["name"], r))
|
|
|
|
# Print clean accuracy row
|
|
print(f"\n{'Clean Accuracy:':<20}", end="")
|
|
for name, r in models_to_show:
|
|
short_name = name[:18] if len(name) > 18 else name
|
|
print(f"{short_name:>18}", end="")
|
|
print()
|
|
print(f"{'':<20}", end="")
|
|
for _, r in models_to_show:
|
|
print(f"{r['clean_acc']:>17.1f}%", end="")
|
|
print()
|
|
|
|
# FGSM table - Model Accuracy (defender's view)
|
|
print(f"\n{'FGSM Model Accuracy (defender):'}")
|
|
print("-" * (20 + 18 * len(models_to_show)))
|
|
print(f"{'Epsilon':<20}", end="")
|
|
for name, _ in models_to_show:
|
|
short_name = name[:18] if len(name) > 18 else name
|
|
print(f"{short_name:>18}", end="")
|
|
print()
|
|
print("-" * (20 + 18 * len(models_to_show)))
|
|
|
|
for eps in epsilon_spread:
|
|
marker = " *" if eps == PRIMARY_EPSILON else ""
|
|
print(f"{eps:<20.2f}", end="")
|
|
for _, r in models_to_show:
|
|
print(f"{r['fgsm_results'][eps]:>17.1f}%", end="")
|
|
print(marker)
|
|
|
|
# FGSM table - Attack Success (attacker's view)
|
|
print(f"\n{'FGSM Attack Success (attacker):'}")
|
|
print("-" * (20 + 18 * len(models_to_show)))
|
|
print(f"{'Epsilon':<20}", end="")
|
|
for name, _ in models_to_show:
|
|
short_name = name[:18] if len(name) > 18 else name
|
|
print(f"{short_name:>18}", end="")
|
|
print()
|
|
print("-" * (20 + 18 * len(models_to_show)))
|
|
|
|
for eps in epsilon_spread:
|
|
marker = " *" if eps == PRIMARY_EPSILON else ""
|
|
print(f"{eps:<20.2f}", end="")
|
|
for _, r in models_to_show:
|
|
attack_success = 100.0 - r["fgsm_results"][eps]
|
|
print(f"{attack_success:>17.1f}%", end="")
|
|
print(marker)
|
|
|
|
# I-FGSM table - Model Accuracy (defender's view)
|
|
print(f"\n{'I-FGSM Model Accuracy (defender):'}")
|
|
print("-" * (20 + 18 * len(models_to_show)))
|
|
print(f"{'Epsilon':<20}", end="")
|
|
for name, _ in models_to_show:
|
|
short_name = name[:18] if len(name) > 18 else name
|
|
print(f"{short_name:>18}", end="")
|
|
print()
|
|
print("-" * (20 + 18 * len(models_to_show)))
|
|
|
|
for eps in epsilon_spread:
|
|
marker = " *" if eps == PRIMARY_EPSILON else ""
|
|
print(f"{eps:<20.2f}", end="")
|
|
for _, r in models_to_show:
|
|
print(f"{r['ifgsm_results'][eps]:>17.1f}%", end="")
|
|
print(marker)
|
|
|
|
# I-FGSM table - Attack Success (attacker's view)
|
|
print(f"\n{'I-FGSM Attack Success (attacker):'}")
|
|
print("-" * (20 + 18 * len(models_to_show)))
|
|
print(f"{'Epsilon':<20}", end="")
|
|
for name, _ in models_to_show:
|
|
short_name = name[:18] if len(name) > 18 else name
|
|
print(f"{short_name:>18}", end="")
|
|
print()
|
|
print("-" * (20 + 18 * len(models_to_show)))
|
|
|
|
for eps in epsilon_spread:
|
|
marker = " *" if eps == PRIMARY_EPSILON else ""
|
|
print(f"{eps:<20.2f}", end="")
|
|
for _, r in models_to_show:
|
|
attack_success = 100.0 - r["ifgsm_results"][eps]
|
|
print(f"{attack_success:>17.1f}%", end="")
|
|
print(marker)
|
|
|
|
print("-" * (20 + 18 * len(models_to_show)))
|
|
print("* Primary epsilon for summary metrics")
|
|
|
|
# Improvement summary at primary epsilon (if baseline exists)
|
|
if baseline:
|
|
print(f"\nImprovement over baseline at ε={PRIMARY_EPSILON}:")
|
|
base_fgsm = baseline["fgsm_results"][PRIMARY_EPSILON]
|
|
base_ifgsm = baseline["ifgsm_results"][PRIMARY_EPSILON]
|
|
for r in all_results:
|
|
fgsm_d = r["fgsm_results"][PRIMARY_EPSILON] - base_fgsm
|
|
ifgsm_d = r["ifgsm_results"][PRIMARY_EPSILON] - base_ifgsm
|
|
print(
|
|
f" {r['name']}: FGSM {'+' if fgsm_d >= 0 else ''}{fgsm_d:.1f}%, I-FGSM {'+' if ifgsm_d >= 0 else ''}{ifgsm_d:.1f}%"
|
|
)
|
|
|
|
print("=" * 90)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Main
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="Evaluate adversarial robustness of trained models",
|
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
|
epilog="""
|
|
Examples:
|
|
python evaluate_robustness.py --model-path robust_model.pth
|
|
python evaluate_robustness.py --model-path models/
|
|
python evaluate_robustness.py --model-path . --compare
|
|
""",
|
|
)
|
|
parser.add_argument(
|
|
"--model-path",
|
|
type=str,
|
|
required=True,
|
|
help="Path to model file (.safetensors) or directory containing models",
|
|
)
|
|
parser.add_argument(
|
|
"--compare",
|
|
action="store_true",
|
|
help="Compare against baseline_model.safetensors",
|
|
)
|
|
parser.add_argument(
|
|
"--show-failures",
|
|
action="store_true",
|
|
help="Show details of misclassified samples",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
# Import torch after argparse for instant --help
|
|
import torch
|
|
|
|
model_path = Path(args.model_path)
|
|
if not model_path.exists():
|
|
print(f"Error: Path not found: {model_path}")
|
|
return 1
|
|
|
|
# Find models to evaluate
|
|
models = find_models(model_path)
|
|
if not models:
|
|
print(f"Error: No .safetensors model files found in {model_path}")
|
|
return 1
|
|
|
|
# Find adversarial examples
|
|
adv_path = find_adv_examples(model_path)
|
|
if not adv_path:
|
|
print("Error: Could not find adv_examples.safetensors")
|
|
return 1
|
|
|
|
# Find baseline if comparing
|
|
baseline_path = None
|
|
if args.compare:
|
|
baseline_path = find_baseline(model_path)
|
|
if not baseline_path:
|
|
print(
|
|
"Warning: --compare specified but baseline_model.safetensors not found"
|
|
)
|
|
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
print(f"Using device: {device}")
|
|
|
|
# Load adversarial examples
|
|
print(f"Loading adversarial examples from: {adv_path.name}")
|
|
adv_data = load_adversarial_examples(adv_path)
|
|
|
|
num_samples = adv_data["clean_images"].size(0)
|
|
has_spread = "epsilon_spread" in adv_data
|
|
print(f" {num_samples} test samples")
|
|
if has_spread:
|
|
print(f" Epsilon spread: {adv_data['epsilon_spread']}")
|
|
else:
|
|
print(f" Single epsilon: {adv_data['epsilon']}")
|
|
|
|
# Evaluate baseline first if comparing
|
|
baseline_results = None
|
|
if baseline_path:
|
|
print(f"\nEvaluating baseline...")
|
|
baseline_results = evaluate_model_full(baseline_path, adv_data, device)
|
|
print_model_results(baseline_results)
|
|
|
|
# Evaluate each model
|
|
all_results = []
|
|
for mp in models:
|
|
print(f"\nEvaluating {mp.name}...")
|
|
results = evaluate_model_full(mp, adv_data, device)
|
|
all_results.append(results)
|
|
print_model_results(results, baseline_results)
|
|
|
|
# Show failures if requested
|
|
if args.show_failures:
|
|
eps_spread = results["epsilon_spread"]
|
|
fgsm_by_eps = adv_data.get(
|
|
"fgsm_by_epsilon", {adv_data["epsilon"]: adv_data.get("fgsm_images")}
|
|
)
|
|
ifgsm_by_eps = adv_data.get(
|
|
"ifgsm_by_epsilon", {adv_data["epsilon"]: adv_data.get("ifgsm_images")}
|
|
)
|
|
|
|
failures = get_misclassified_samples(
|
|
results["model"],
|
|
ifgsm_by_eps[PRIMARY_EPSILON],
|
|
adv_data["clean_labels"],
|
|
device,
|
|
)
|
|
if failures:
|
|
print(f"\nMisclassified samples (I-FGSM ε={PRIMARY_EPSILON}):")
|
|
for f in failures:
|
|
print(
|
|
f" Sample {f['index']}: {f['true_label']} → {f['predicted']} (conf={f['confidence']:.2f})"
|
|
)
|
|
|
|
# Print comparison summary if multiple models or baseline comparison
|
|
if len(all_results) > 1 or baseline_results:
|
|
print_comparison_summary(all_results, baseline_results)
|
|
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
exit(main())
|