jjanella/AI-Red-Teaming-CSCD94
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

added material

Browse Source
This commit is contained in:
Jeremy Janella
2026-05-09 23:21:13 -04:00
commit 75faa5c410
369 changed files with 3468754 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
defense/adversarial training/adv_examples.safetensors
BIN
View File
Binary file not shown.
defense/adversarial training/adversarial training.ipynb
+714
View File
@@ -0,0 +1,714 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"id": "63b9ead7-f9e3-4edd-855c-4ec0d9b59e15",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz\n",
"Failed to download (trying next):\n",
"HTTP Error 404: Not Found\n",
"\n",
"Downloading https://ossci-datasets.s3.amazonaws.com/mnist/train-images-idx3-ubyte.gz\n",
"Downloading https://ossci-datasets.s3.amazonaws.com/mnist/train-images-idx3-ubyte.gz to ./data/MNIST/raw/train-images-idx3-ubyte.gz\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"100%|████████████████████████████████████████████████████████████████████████████████████████████████████████| 9.91M/9.91M [00:00<00:00, 51.7MB/s]\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Extracting ./data/MNIST/raw/train-images-idx3-ubyte.gz to ./data/MNIST/raw\n",
"\n",
"Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz\n",
"Failed to download (trying next):\n",
"HTTP Error 404: Not Found\n",
"\n",
"Downloading https://ossci-datasets.s3.amazonaws.com/mnist/train-labels-idx1-ubyte.gz\n",
"Downloading https://ossci-datasets.s3.amazonaws.com/mnist/train-labels-idx1-ubyte.gz to ./data/MNIST/raw/train-labels-idx1-ubyte.gz\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"100%|████████████████████████████████████████████████████████████████████████████████████████████████████████| 28.9k/28.9k [00:00<00:00, 1.38MB/s]"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Extracting ./data/MNIST/raw/train-labels-idx1-ubyte.gz to ./data/MNIST/raw\n",
"\n",
"Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz\n",
"Failed to download (trying next):\n",
"HTTP Error 404: Not Found\n",
"\n",
"Downloading https://ossci-datasets.s3.amazonaws.com/mnist/t10k-images-idx3-ubyte.gz\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Downloading https://ossci-datasets.s3.amazonaws.com/mnist/t10k-images-idx3-ubyte.gz to ./data/MNIST/raw/t10k-images-idx3-ubyte.gz\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"100%|████████████████████████████████████████████████████████████████████████████████████████████████████████| 1.65M/1.65M [00:00<00:00, 11.4MB/s]\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Extracting ./data/MNIST/raw/t10k-images-idx3-ubyte.gz to ./data/MNIST/raw\n",
"\n",
"Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz\n",
"Failed to download (trying next):\n",
"HTTP Error 404: Not Found\n",
"\n",
"Downloading https://ossci-datasets.s3.amazonaws.com/mnist/t10k-labels-idx1-ubyte.gz\n",
"Downloading https://ossci-datasets.s3.amazonaws.com/mnist/t10k-labels-idx1-ubyte.gz to ./data/MNIST/raw/t10k-labels-idx1-ubyte.gz\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"100%|████████████████████████████████████████████████████████████████████████████████████████████████████████| 4.54k/4.54k [00:00<00:00, 5.35MB/s]\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Extracting ./data/MNIST/raw/t10k-labels-idx1-ubyte.gz to ./data/MNIST/raw\n",
"\n",
"Epoch 1/10: Avg Loss = 0.3898, Test Accuracy = 96.67%\n",
"Epoch 2/10: Avg Loss = 0.1010, Test Accuracy = 97.92%\n",
"Epoch 3/10: Avg Loss = 0.0693, Test Accuracy = 98.43%\n",
"Epoch 4/10: Avg Loss = 0.0523, Test Accuracy = 98.18%\n",
"Epoch 5/10: Avg Loss = 0.0448, Test Accuracy = 98.82%\n",
"Epoch 6/10: Avg Loss = 0.0368, Test Accuracy = 98.56%\n",
"Epoch 7/10: Avg Loss = 0.0327, Test Accuracy = 98.95%\n",
"Epoch 8/10: Avg Loss = 0.0280, Test Accuracy = 99.02%\n",
"Epoch 9/10: Avg Loss = 0.0254, Test Accuracy = 98.89%\n",
"Epoch 10/10: Avg Loss = 0.0211, Test Accuracy = 99.02%\n"
]
}
],
"source": [
"import json\n",
"import numpy as np\n",
"import torch\n",
"import torch.nn as nn\n",
"import torch.nn.functional as F\n",
"import torch.optim as optim\n",
"from safetensors.torch import save_file\n",
"from tqdm import tqdm\n",
"\n",
"from htb_ai_library import (\n",
" set_reproducibility,\n",
" get_mnist_loaders,\n",
" evaluate_accuracy,\n",
" train_model,\n",
")\n",
"\n",
"MNIST_MEAN = 0.1307\n",
"MNIST_STD = 0.3081\n",
"EPSILON = 0.3\n",
"EPSILON_SPREAD = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]\n",
"I_FGSM_STEPS = 10\n",
"\n",
"\n",
"def get_device():\n",
" \"\"\"Get the best available device.\"\"\"\n",
" if torch.cuda.is_available():\n",
" return torch.device(\"cuda\")\n",
" return torch.device(\"cpu\")\n",
"\n",
"def save_adversarial_examples(data, path):\n",
" \"\"\"Save adversarial examples to safetensors format.\"\"\"\n",
" # We don't store fgsm_images/ifgsm_images separately since they reference\n",
" # fgsm_by_epsilon[epsilon] and would cause memory sharing errors in safetensors\n",
" tensors = {\n",
" 'clean_images': data['clean_images'],\n",
" 'clean_labels': data['clean_labels'].long(),\n",
" }\n",
"\n",
" for eps in data['epsilon_spread']:\n",
" key = f\"{eps:.1f}\"\n",
" tensors[f'fgsm_eps_{key}'] = data['fgsm_by_epsilon'][eps]\n",
" tensors[f'ifgsm_eps_{key}'] = data['ifgsm_by_epsilon'][eps]\n",
"\n",
" metadata = {\n",
" 'epsilon': str(data['epsilon']),\n",
" 'epsilon_spread': json.dumps(data['epsilon_spread']),\n",
" }\n",
"\n",
" save_file(tensors, path, metadata=metadata)\n",
"\n",
"class LeNet5(nn.Module):\n",
" def __init__(self):\n",
" super().__init__()\n",
" self.conv1 = nn.Conv2d(1, 6, kernel_size=5, padding=2)\n",
" self.conv2 = nn.Conv2d(6, 16, kernel_size=5)\n",
" self.fc1 = nn.Linear(16 * 5 * 5, 120)\n",
" self.fc2 = nn.Linear(120, 84)\n",
" self.fc3 = nn.Linear(84, 10)\n",
"\n",
" def forward(self, x):\n",
" x = F.max_pool2d(F.relu(self.conv1(x)), 2)\n",
" x = F.max_pool2d(F.relu(self.conv2(x)), 2)\n",
" x = x.view(-1, 16 * 5 * 5)\n",
" x = F.relu(self.fc1(x))\n",
" x = F.relu(self.fc2(x))\n",
" x = self.fc3(x)\n",
" return x\n",
"\n",
"def fgsm_attack(model, images, labels, epsilon):\n",
" images_copy = images.clone().detach().requires_grad_(True)\n",
" outputs = model(images_copy)\n",
" loss = F.cross_entropy(outputs, labels)\n",
" model.zero_grad()\n",
" loss.backward()\n",
" grad_sign = images_copy.grad.sign()\n",
" adv_images = images_copy + epsilon * grad_sign\n",
" min_val = (0 - MNIST_MEAN) / MNIST_STD\n",
" max_val = (1 - MNIST_MEAN) / MNIST_STD\n",
" adv_images = torch.clamp(adv_images, min_val, max_val)\n",
" return adv_images.detach()\n",
"\n",
"def i_fgsm_attack(model, images, labels, epsilon, steps=I_FGSM_STEPS):\n",
" \"\"\"Generate I-FGSM (Iterative FGSM) adversarial examples.\"\"\"\n",
" alpha = epsilon / steps # Step size per iteration\n",
" min_val = (0 - MNIST_MEAN) / MNIST_STD\n",
" max_val = (1 - MNIST_MEAN) / MNIST_STD\n",
"\n",
" adv_images = images.clone().detach()\n",
" original_images = images.clone().detach()\n",
"\n",
" for _ in range(steps):\n",
" adv_images.requires_grad = True\n",
" outputs = model(adv_images)\n",
" loss = F.cross_entropy(outputs, labels)\n",
" model.zero_grad()\n",
" loss.backward()\n",
"\n",
" grad_sign = adv_images.grad.sign()\n",
" adv_images = adv_images.detach() + alpha * grad_sign\n",
"\n",
" # Project back to epsilon-ball around original\n",
" perturbation = adv_images - original_images\n",
" perturbation = torch.clamp(perturbation, -epsilon, epsilon)\n",
" adv_images = original_images + perturbation\n",
"\n",
" # Clamp to valid range\n",
" adv_images = torch.clamp(adv_images, min_val, max_val)\n",
"\n",
" return adv_images.detach()\n",
"\n",
"def evaluate_adversarial_accuracy(model, loader, device, epsilon, num_batches=None):\n",
" \"\"\"Evaluate accuracy under FGSM attack.\"\"\"\n",
" model.eval()\n",
" correct = 0\n",
" total = 0\n",
"\n",
" for i, (images, labels) in enumerate(loader):\n",
" if num_batches is not None and i >= num_batches:\n",
" break\n",
"\n",
" images, labels = images.to(device), labels.to(device)\n",
"\n",
" # Generate adversarial examples (need gradients, so briefly enable train mode)\n",
" model.train()\n",
" adv_images = fgsm_attack(model, images, labels, epsilon)\n",
" model.eval()\n",
"\n",
" with torch.no_grad():\n",
" outputs = model(adv_images)\n",
" _, predicted = outputs.max(1)\n",
" total += labels.size(0)\n",
" correct += predicted.eq(labels).sum().item()\n",
"\n",
" return 100.0 * correct / total\n",
"\n",
"train_loader, test_loader = get_mnist_loaders(batch_size=128, data_dir=\"./data\")\n",
"baseline_model = LeNet5()\n",
"baseline_model = train_model(\n",
" baseline_model,\n",
" train_loader,\n",
" test_loader,\n",
" device=get_device(),\n",
" epochs=10,\n",
" learning_rate=0.001,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "b1dfab4c-0326-4093-b3f9-71489525129e",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Using device: cpu\n",
"Train batches: 469, Test batches: 79\n"
]
}
],
"source": [
"def generate_adversarial_examples(model, test_loader, device, num_samples=500):\n",
" \"\"\"Generate adversarial examples across multiple epsilon values.\"\"\"\n",
" model.eval()\n",
"\n",
" # Collect clean samples\n",
" clean_images_list = []\n",
" clean_labels_list = []\n",
" collected = 0\n",
"\n",
" print(f\"Collecting {num_samples} clean samples...\")\n",
" for images, labels in test_loader:\n",
" if collected >= num_samples:\n",
" break\n",
" batch_size = min(images.size(0), num_samples - collected)\n",
" clean_images_list.append(images[:batch_size])\n",
" clean_labels_list.append(labels[:batch_size])\n",
" collected += batch_size\n",
"\n",
" clean_images = torch.cat(clean_images_list, dim=0)\n",
" clean_labels = torch.cat(clean_labels_list, dim=0)\n",
"\n",
" # Generate adversarial examples at each epsilon\n",
" fgsm_by_epsilon = {}\n",
" ifgsm_by_epsilon = {}\n",
"\n",
" print(f\"\\nGenerating adversarial examples across epsilon spread: {EPSILON_SPREAD}\")\n",
"\n",
" for eps in EPSILON_SPREAD:\n",
" print(f\"\\n Generating at epsilon={eps}...\")\n",
" fgsm_images_list = []\n",
" ifgsm_images_list = []\n",
"\n",
" batch_size = 128\n",
" pbar = tqdm(total=num_samples, desc=f\" eps={eps}\")\n",
"\n",
" for i in range(0, num_samples, batch_size):\n",
" end_idx = min(i + batch_size, num_samples)\n",
" images = clean_images[i:end_idx].to(device)\n",
" labels = clean_labels[i:end_idx].to(device)\n",
"\n",
" model.train() # Need train mode for gradient computation\n",
" fgsm_images = fgsm_attack(model, images, labels, eps)\n",
" ifgsm_images = i_fgsm_attack(model, images, labels, eps)\n",
" model.eval()\n",
"\n",
" fgsm_images_list.append(fgsm_images.cpu())\n",
" ifgsm_images_list.append(ifgsm_images.cpu())\n",
" pbar.update(end_idx - i)\n",
"\n",
" pbar.close()\n",
"\n",
" fgsm_by_epsilon[eps] = torch.cat(fgsm_images_list, dim=0)\n",
" ifgsm_by_epsilon[eps] = torch.cat(ifgsm_images_list, dim=0)\n",
"\n",
" return {\n",
" 'clean_images': clean_images,\n",
" 'clean_labels': clean_labels,\n",
" 'fgsm_images': fgsm_by_epsilon[EPSILON],\n",
" 'ifgsm_images': ifgsm_by_epsilon[EPSILON],\n",
" 'epsilon': EPSILON,\n",
" 'epsilon_spread': EPSILON_SPREAD,\n",
" 'fgsm_by_epsilon': fgsm_by_epsilon,\n",
" 'ifgsm_by_epsilon': ifgsm_by_epsilon\n",
" }\n",
"\n",
"def train_adversarial(model, train_loader, test_loader, device,\n",
" epochs=25, lr=0.001, epsilon=EPSILON):\n",
" \"\"\"Train model with adversarial training using epsilon spread.\"\"\"\n",
" model.to(device)\n",
"\n",
" optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4)\n",
" scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)\n",
" criterion = nn.CrossEntropyLoss()\n",
"\n",
" for epoch in range(epochs):\n",
" model.train()\n",
" total_loss = 0.0\n",
" correct = 0\n",
" total = 0\n",
"\n",
" pbar = tqdm(train_loader, desc=f\"Epoch {epoch+1}/{epochs}\")\n",
" for images, labels in pbar:\n",
" images, labels = images.to(device), labels.to(device)\n",
"\n",
" batch_epsilon = np.random.choice(EPSILON_SPREAD)\n",
" adv_images = fgsm_attack(model, images, labels, batch_epsilon)\n",
"\n",
" combined_images = torch.cat([images, adv_images], dim=0)\n",
" combined_labels = torch.cat([labels, labels], dim=0)\n",
"\n",
" perm = torch.randperm(combined_images.size(0))\n",
" combined_images = combined_images[perm]\n",
" combined_labels = combined_labels[perm]\n",
"\n",
" optimizer.zero_grad()\n",
" outputs = model(combined_images)\n",
" loss = criterion(outputs, combined_labels)\n",
" loss.backward()\n",
" torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)\n",
" optimizer.step()\n",
"\n",
" total_loss += loss.item()\n",
" _, predicted = outputs.max(1)\n",
" total += combined_labels.size(0)\n",
" correct += predicted.eq(combined_labels).sum().item()\n",
"\n",
" pbar.set_postfix({\n",
" 'loss': f'{total_loss / (pbar.n + 1):.4f}',\n",
" 'acc': f'{100.0 * correct / total:.2f}%'\n",
" })\n",
"\n",
" scheduler.step()\n",
"\n",
" if (epoch + 1) % 5 == 0 or epoch == epochs - 1:\n",
" clean_acc = evaluate_accuracy(model, test_loader, device)\n",
" adv_acc = evaluate_adversarial_accuracy(\n",
" model, test_loader, device, epsilon, num_batches=20\n",
" )\n",
" print(f\"\\n Epoch {epoch+1}: Clean={clean_acc:.1f}%, Robust={adv_acc:.1f}%\")\n",
"\n",
" return model\n",
"\n",
"set_reproducibility(1337)\n",
"device = get_device()\n",
"print(f\"Using device: {device}\")\n",
"\n",
"train_loader, test_loader = get_mnist_loaders(normalize=True)\n",
"print(f\"Train batches: {len(train_loader)}, Test batches: {len(test_loader)}\")"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "658b5b4c-74f9-44bf-a959-8ea7ab27b7c5",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 1/10: Avg Loss = 0.3098, Test Accuracy = 97.53%\n",
"Epoch 2/10: Avg Loss = 0.0809, Test Accuracy = 98.48%\n",
"Epoch 3/10: Avg Loss = 0.0539, Test Accuracy = 98.63%\n",
"Epoch 4/10: Avg Loss = 0.0431, Test Accuracy = 98.49%\n",
"Epoch 5/10: Avg Loss = 0.0348, Test Accuracy = 98.78%\n",
"Epoch 6/10: Avg Loss = 0.0295, Test Accuracy = 98.71%\n",
"Epoch 7/10: Avg Loss = 0.0257, Test Accuracy = 98.78%\n",
"Epoch 8/10: Avg Loss = 0.0204, Test Accuracy = 98.84%\n",
"Epoch 9/10: Avg Loss = 0.0167, Test Accuracy = 98.96%\n",
"Epoch 10/10: Avg Loss = 0.0151, Test Accuracy = 98.88%\n",
"Baseline Model - Robust: 74.8%\n"
]
}
],
"source": [
"baseline_model = LeNet5()\n",
"baseline_model = train_model(baseline_model, train_loader, test_loader,\n",
" device=device, epochs=10, learning_rate=0.001)\n",
"\n",
"adv_acc = evaluate_adversarial_accuracy(baseline_model, test_loader, device, EPSILON)\n",
"print(f\"Baseline Model - Robust: {adv_acc:.1f}%\")\n",
"save_file(baseline_model.state_dict(), \"baseline_model.safetensors\")"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "eca344cf-f1b2-43d3-b563-8a618fd8f97a",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Collecting 500 clean samples...\n",
"\n",
"Generating adversarial examples across epsilon spread: [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]\n",
"\n",
" Generating at epsilon=0.1...\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
" eps=0.1: 100%|██████████████████████████████████████████████████████████████████████████████████████████████| 500/500 [00:00<00:00, 3589.91it/s]\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n",
" Generating at epsilon=0.2...\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
" eps=0.2: 100%|██████████████████████████████████████████████████████████████████████████████████████████████| 500/500 [00:00<00:00, 3783.98it/s]\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n",
" Generating at epsilon=0.3...\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
" eps=0.3: 100%|██████████████████████████████████████████████████████████████████████████████████████████████| 500/500 [00:00<00:00, 4337.29it/s]\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n",
" Generating at epsilon=0.4...\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
" eps=0.4: 100%|██████████████████████████████████████████████████████████████████████████████████████████████| 500/500 [00:00<00:00, 4218.09it/s]\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n",
" Generating at epsilon=0.5...\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
" eps=0.5: 100%|██████████████████████████████████████████████████████████████████████████████████████████████| 500/500 [00:00<00:00, 2903.09it/s]\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n",
" Generating at epsilon=0.6...\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
" eps=0.6: 100%|██████████████████████████████████████████████████████████████████████████████████████████████| 500/500 [00:00<00:00, 2815.23it/s]\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n",
" Generating at epsilon=0.7...\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
" eps=0.7: 100%|██████████████████████████████████████████████████████████████████████████████████████████████| 500/500 [00:00<00:00, 2339.08it/s]\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n",
" Generating at epsilon=0.8...\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
" eps=0.8: 100%|██████████████████████████████████████████████████████████████████████████████████████████████| 500/500 [00:00<00:00, 2391.63it/s]\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n",
" Generating at epsilon=0.9...\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
" eps=0.9: 100%|██████████████████████████████████████████████████████████████████████████████████████████████| 500/500 [00:00<00:00, 2493.57it/s]\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n",
" Generating at epsilon=1.0...\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
" eps=1.0: 100%|██████████████████████████████████████████████████████████████████████████████████████████████| 500/500 [00:00<00:00, 2372.05it/s]\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Adversarial examples saved to adv_examples.safetensors\n",
"Before training - Clean: 10.2%, Robust: 2.4%\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"Epoch 1/10: 100%|██████████████████████████████████████████████████████████████████████| 469/469 [00:06<00:00, 70.61it/s, loss=0.7535, acc=75.68%]\n",
"Epoch 2/10: 100%|██████████████████████████████████████████████████████████████████████| 469/469 [00:06<00:00, 73.95it/s, loss=0.3464, acc=88.43%]\n",
"Epoch 3/10: 100%|██████████████████████████████████████████████████████████████████████| 469/469 [00:06<00:00, 71.96it/s, loss=0.2837, acc=90.47%]\n",
"Epoch 4/10: 100%|██████████████████████████████████████████████████████████████████████| 469/469 [00:06<00:00, 72.84it/s, loss=0.2241, acc=92.43%]\n",
"Epoch 5/10: 100%|██████████████████████████████████████████████████████████████████████| 469/469 [00:06<00:00, 72.62it/s, loss=0.1920, acc=93.46%]\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n",
" Epoch 5: Clean=98.9%, Robust=92.9%\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"Epoch 6/10: 100%|██████████████████████████████████████████████████████████████████████| 469/469 [00:06<00:00, 71.49it/s, loss=0.1725, acc=94.28%]\n",
"Epoch 7/10: 100%|██████████████████████████████████████████████████████████████████████| 469/469 [00:06<00:00, 72.60it/s, loss=0.1551, acc=94.83%]\n",
"Epoch 8/10: 100%|██████████████████████████████████████████████████████████████████████| 469/469 [00:06<00:00, 73.79it/s, loss=0.1474, acc=95.11%]\n",
"Epoch 9/10: 100%|██████████████████████████████████████████████████████████████████████| 469/469 [00:06<00:00, 73.62it/s, loss=0.1393, acc=95.43%]\n",
"Epoch 10/10: 100%|█████████████████████████████████████████████████████████████████████| 469/469 [00:06<00:00, 72.08it/s, loss=0.1319, acc=95.60%]\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n",
" Epoch 10: Clean=99.0%, Robust=93.8%\n",
"Final - Clean: 99.0%, Robust: 95.7%\n",
"Model saved to robust_model.safetensors\n"
]
}
],
"source": [
"adv_data = generate_adversarial_examples(\n",
" baseline_model, test_loader, device, num_samples=500\n",
")\n",
"\n",
"save_adversarial_examples(adv_data, \"adv_examples.safetensors\")\n",
"print(\"Adversarial examples saved to adv_examples.safetensors\")\n",
"\n",
"model = LeNet5()\n",
"model.to(device)\n",
"\n",
"clean_acc = evaluate_accuracy(model, test_loader, device)\n",
"adv_acc = evaluate_adversarial_accuracy(model, test_loader, device, EPSILON)\n",
"print(f\"Before training - Clean: {clean_acc:.1f}%, Robust: {adv_acc:.1f}%\")\n",
"\n",
"model = train_adversarial(model, train_loader, test_loader, device, epochs=10)\n",
"\n",
"clean_acc = evaluate_accuracy(model, test_loader, device)\n",
"adv_acc = evaluate_adversarial_accuracy(model, test_loader, device, EPSILON)\n",
"print(f\"Final - Clean: {clean_acc:.1f}%, Robust: {adv_acc:.1f}%\")\n",
"\n",
"save_file(model.state_dict(), \"robust_model.safetensors\")\n",
"print(\"Model saved to robust_model.safetensors\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "2b8045fa-9e0b-468c-b46c-0dedbab41f61",
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.6"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
defense/adversarial training/adversarial_training.zip
BIN
View File
Binary file not shown.
defense/adversarial training/baseline_model.safetensors
BIN
View File
Binary file not shown.
defense/adversarial training/data/MNIST/raw/t10k-images-idx3-ubyte
BIN
View File
Binary file not shown.
defense/adversarial training/data/MNIST/raw/t10k-images-idx3-ubyte.gz
BIN
View File
Binary file not shown.
defense/adversarial training/data/MNIST/raw/t10k-labels-idx1-ubyte
BIN
View File
Binary file not shown.
defense/adversarial training/data/MNIST/raw/t10k-labels-idx1-ubyte.gz
BIN
View File
Binary file not shown.
defense/adversarial training/data/MNIST/raw/train-images-idx3-ubyte
BIN
View File
Binary file not shown.
defense/adversarial training/data/MNIST/raw/train-images-idx3-ubyte.gz
BIN
View File
Binary file not shown.
defense/adversarial training/data/MNIST/raw/train-labels-idx1-ubyte
BIN
View File
Binary file not shown.
defense/adversarial training/data/MNIST/raw/train-labels-idx1-ubyte.gz
BIN
View File
Binary file not shown.
defense/adversarial training/evaluate_robustness.py
+581
View File
@@ -0,0 +1,581 @@
#!/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())
defense/adversarial training/robust_model.safetensors
BIN
View File
Binary file not shown.