jjanella/AI-Red-Teaming-CSCD94
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
95b1c6bf27226c2758482da10ed62a8def42df39
AI-Red-Teaming-CSCD94/defense/adversarial training/adversarial training.ipynb
T
Jeremy Janella 95b1c6bf27 added material
2026-07-26 22:53:03 -04:00

715 lines
30 KiB
Plaintext
Raw Blame History

{
"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
}
Reference in New Issue View Git Blame Copy Permalink