{ "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 }