jjanella/AI-Red-Teaming-CSCD94
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
main
AI-Red-Teaming-CSCD94/ai-data-attacks/trojan-attack/trojan-attack.ipynb
T
Jeremy Janella 75faa5c410 added material
2026-05-09 23:21:13 -04:00

2668 lines
96 KiB
Plaintext
Raw Permalink Blame History

{
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"id": "30606703-3597-4ae0-9c2a-cb40bbf74f0f",
"metadata": {},
"outputs": [],
"source": [
"import torch\n",
"import torch.nn as nn\n",
"import torch.optim as optim\n",
"import torch.nn.functional as F"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "60ae3dc5-9f16-4396-909b-3e38b77357d4",
"metadata": {},
"outputs": [],
"source": [
"import torchvision\n",
"import torchvision.transforms as transforms\n",
"from torch.utils.data import DataLoader, Dataset\n",
"from torchvision.datasets import ImageFolder"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "3b350220-ace2-453e-a940-4bffbbc0ce01",
"metadata": {},
"outputs": [],
"source": [
"from tqdm.auto import tqdm, trange\n",
"import numpy as np\n",
"import matplotlib.pyplot as plt\n",
"import random\n",
"import copy\n",
"import os\n",
"import pandas as pd\n",
"from PIL import Image\n",
"import requests\n",
"import zipfile\n",
"import shutil"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "ca22daec-f61f-47f6-ab9f-a7be914645da",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Using CPU device.\n",
"Using device: cpu\n"
]
}
],
"source": [
"# Enforce determinism for reproducibility\n",
"torch.backends.cudnn.deterministic = True\n",
"torch.backends.cudnn.benchmark = False\n",
"\n",
"# Device configuration\n",
"if torch.cuda.is_available():\n",
" device = torch.device(\"cuda\")\n",
" print(\"Using CUDA device.\")\n",
"elif torch.backends.mps.is_available():\n",
" device = torch.device(\"mps\")\n",
" print(\"Using MPS device (Apple Silicon GPU).\")\n",
"else:\n",
" device = torch.device(\"cpu\")\n",
" print(\"Using CPU device.\")\n",
"print(f\"Using device: {device}\")\n"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "c959b777-fa15-432f-9914-7a6a442244df",
"metadata": {},
"outputs": [],
"source": [
"# Set random seed for reproducibility\n",
"SEED = 1337\n",
"random.seed(SEED)\n",
"np.random.seed(SEED)\n",
"torch.manual_seed(SEED)\n",
"if torch.cuda.is_available(): # Ensure CUDA seeds are set only if GPU is used\n",
" torch.cuda.manual_seed(SEED)\n",
" torch.cuda.manual_seed_all(SEED) # For multi-GPU setups\n"
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "15139d30-01c7-4585-83ee-8af859206286",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Setup complete.\n"
]
}
],
"source": [
"# Primary Palette\n",
"HTB_GREEN = \"#9fef00\"\n",
"NODE_BLACK = \"#141d2b\"\n",
"HACKER_GREY = \"#a4b1cd\"\n",
"WHITE = \"#ffffff\"\n",
"# Secondary Palette\n",
"AZURE = \"#0086ff\"\n",
"NUGGET_YELLOW = \"#ffaf00\"\n",
"MALWARE_RED = \"#ff3e3e\"\n",
"VIVID_PURPLE = \"#9f00ff\"\n",
"AQUAMARINE = \"#2ee7b6\"\n",
"# Matplotlib Style Settings\n",
"plt.style.use(\"seaborn-v0_8-darkgrid\")\n",
"plt.rcParams.update(\n",
" {\n",
" \"figure.facecolor\": NODE_BLACK,\n",
" \"figure.edgecolor\": NODE_BLACK,\n",
" \"axes.facecolor\": NODE_BLACK,\n",
" \"axes.edgecolor\": HACKER_GREY,\n",
" \"axes.labelcolor\": HACKER_GREY,\n",
" \"axes.titlecolor\": WHITE,\n",
" \"xtick.color\": HACKER_GREY,\n",
" \"ytick.color\": HACKER_GREY,\n",
" \"grid.color\": HACKER_GREY,\n",
" \"grid.alpha\": 0.1,\n",
" \"legend.facecolor\": NODE_BLACK,\n",
" \"legend.edgecolor\": HACKER_GREY,\n",
" \"legend.labelcolor\": HACKER_GREY,\n",
" \"text.color\": HACKER_GREY,\n",
" }\n",
")\n",
"\n",
"print(\"Setup complete.\")\n"
]
},
{
"cell_type": "code",
"execution_count": 7,
"id": "5ef3af09-4897-4757-9e64-9cbff7cd320c",
"metadata": {},
"outputs": [],
"source": [
"GTSRB_CLASS_NAMES = {\n",
" 0: \"Speed limit (20km/h)\",\n",
" 1: \"Speed limit (30km/h)\",\n",
" 2: \"Speed limit (50km/h)\",\n",
" 3: \"Speed limit (60km/h)\",\n",
" 4: \"Speed limit (70km/h)\",\n",
" 5: \"Speed limit (80km/h)\",\n",
" 6: \"End of speed limit (80km/h)\",\n",
" 7: \"Speed limit (100km/h)\",\n",
" 8: \"Speed limit (120km/h)\",\n",
" 9: \"No passing\",\n",
" 10: \"No passing for veh over 3.5 tons\",\n",
" 11: \"Right-of-way at next intersection\",\n",
" 12: \"Priority road\",\n",
" 13: \"Yield\",\n",
" 14: \"Stop\",\n",
" 15: \"No vehicles\",\n",
" 16: \"Veh > 3.5 tons prohibited\",\n",
" 17: \"No entry\",\n",
" 18: \"General caution\",\n",
" 19: \"Dangerous curve left\",\n",
" 20: \"Dangerous curve right\",\n",
" 21: \"Double curve\",\n",
" 22: \"Bumpy road\",\n",
" 23: \"Slippery road\",\n",
" 24: \"Road narrows on the right\",\n",
" 25: \"Road work\",\n",
" 26: \"Traffic signals\",\n",
" 27: \"Pedestrians\",\n",
" 28: \"Children crossing\",\n",
" 29: \"Bicycles crossing\",\n",
" 30: \"Beware of ice/snow\",\n",
" 31: \"Wild animals crossing\",\n",
" 32: \"End speed/pass limits\",\n",
" 33: \"Turn right ahead\",\n",
" 34: \"Turn left ahead\",\n",
" 35: \"Ahead only\",\n",
" 36: \"Go straight or right\",\n",
" 37: \"Go straight or left\",\n",
" 38: \"Keep right\",\n",
" 39: \"Keep left\",\n",
" 40: \"Roundabout mandatory\",\n",
" 41: \"End of no passing\",\n",
" 42: \"End no passing veh > 3.5 tons\",\n",
"}\n",
"NUM_CLASSES_GTSRB = len(GTSRB_CLASS_NAMES) # Should be 43\n",
"\n",
"\n",
"def get_gtsrb_class_name(class_id):\n",
" \"\"\"\n",
" Retrieves the human-readable name for a given GTSRB class ID.\n",
"\n",
" Args:\n",
" class_id (int): The numeric class ID (0-42).\n",
"\n",
" Returns:\n",
" str: The corresponding class name or an 'Unknown Class' string.\n",
" \"\"\"\n",
" return GTSRB_CLASS_NAMES.get(class_id, f\"Unknown Class {class_id}\")"
]
},
{
"cell_type": "code",
"execution_count": 8,
"id": "8e6418ec-e06f-4f4d-b768-f4359865821c",
"metadata": {},
"outputs": [],
"source": [
"# Dataset Root Directory\n",
"DATASET_ROOT = \"./GTSRB\"\n",
"\n",
"# URLs for the GTSRB dataset components\n",
"DATASET_URL = \"https://academy.hackthebox.com/storage/resources/GTSRB.zip\"\n",
"DOWNLOAD_DIR = \"./gtsrb_downloads\" # Temporary download location\n",
"\n",
"\n",
"def download_file(url, dest_folder, filename):\n",
" \"\"\"\n",
" Downloads a file from a URL to a specified destination.\n",
"\n",
" Args:\n",
" url (str): The URL of the file to download.\n",
" dest_folder (str): The directory to save the downloaded file.\n",
" filename (str): The name to save the file as.\n",
"\n",
" Returns:\n",
" str or None: The full path to the downloaded file, or None if download failed.\n",
" \"\"\"\n",
" filepath = os.path.join(dest_folder, filename)\n",
" if os.path.exists(filepath):\n",
" print(f\"File '{filename}' already exists in {dest_folder}. Skipping download.\")\n",
" return filepath\n",
" print(f\"Downloading {filename} from {url}...\")\n",
" try:\n",
" response = requests.get(url, stream=True)\n",
" response.raise_for_status() # Raise an exception for bad status codes\n",
" os.makedirs(dest_folder, exist_ok=True)\n",
" with open(filepath, \"wb\") as f:\n",
" for chunk in response.iter_content(chunk_size=8192):\n",
" f.write(chunk)\n",
" print(f\"Successfully downloaded {filename}.\")\n",
" return filepath\n",
" except requests.exceptions.RequestException as e:\n",
" print(f\"Error downloading {url}: {e}\")\n",
" return None\n",
"\n",
"\n",
"def extract_zip(zip_filepath, extract_to):\n",
" \"\"\"\n",
" Extracts the contents of a zip file to a specified directory.\n",
"\n",
" Args:\n",
" zip_filepath (str): The path to the zip file.\n",
" extract_to (str): The directory where contents should be extracted.\n",
"\n",
" Returns:\n",
" bool: True if extraction was successful, False otherwise.\n",
" \"\"\"\n",
" print(f\"Extracting '{os.path.basename(zip_filepath)}' to {extract_to}...\")\n",
" try:\n",
" with zipfile.ZipFile(zip_filepath, \"r\") as zip_ref:\n",
" zip_ref.extractall(extract_to)\n",
" print(f\"Successfully extracted '{os.path.basename(zip_filepath)}'.\")\n",
" return True\n",
" except zipfile.BadZipFile:\n",
" print(\n",
" f\"Error: Failed to extract '{os.path.basename(zip_filepath)}'. File might be corrupted or not a zip file.\"\n",
" )\n",
" return False\n",
" except Exception as e:\n",
" print(f\"An unexpected error occurred during extraction: {e}\")\n",
" return False\n"
]
},
{
"cell_type": "code",
"execution_count": 9,
"id": "3ff62bdc-0e76-4bdb-bcc7-1f5d919df155",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"GTSRB dataset found and seems complete in './GTSRB'. Skipping download.\n"
]
}
],
"source": [
"# Define expected paths within DATASET_ROOT\n",
"train_dir = os.path.join(DATASET_ROOT, \"Final_Training\", \"Images\")\n",
"test_img_dir = os.path.join(DATASET_ROOT, \"Final_Test\", \"Images\")\n",
"test_csv_path = os.path.join(DATASET_ROOT, \"GT-final_test.csv\")\n",
"\n",
"# Check if the core dataset components exist\n",
"dataset_ready = (\n",
" os.path.isdir(DATASET_ROOT)\n",
" and os.path.isdir(train_dir)\n",
" and os.path.isdir(test_img_dir) # Check if test dir exists\n",
" and os.path.isfile(test_csv_path) # Check if test csv exists\n",
")\n",
"\n",
"if dataset_ready:\n",
" print(\n",
" f\"GTSRB dataset found and seems complete in '{DATASET_ROOT}'. Skipping download.\"\n",
" )\n",
"else:\n",
" print(\n",
" f\"GTSRB dataset not found or incomplete in '{DATASET_ROOT}'. Attempting download and extraction...\"\n",
" )\n",
" os.makedirs(DATASET_ROOT, exist_ok=True)\n",
" os.makedirs(DOWNLOAD_DIR, exist_ok=True)\n",
"\n",
" # Download files\n",
" dataset_zip_path = download_file(\n",
" DATASET_URL, DOWNLOAD_DIR, \"GTSRB.zip\"\n",
" )\n",
" extraction_ok = True\n",
" # Only extract if download happened and train_dir doesn't already exist\n",
" if dataset_zip_path and not os.path.isdir(train_dir):\n",
" if not extract_zip(dataset_zip_path, DATASET_ROOT):\n",
" extraction_ok = False\n",
" print(\"Error during extraction of training images.\")\n",
" elif not dataset_zip_path and not os.path.isdir(train_dir):\n",
" # If download failed AND train dir doesn't exist, extraction can't happen\n",
" extraction_ok = False\n",
" print(\"Training images download failed or skipped, cannot proceed with extraction.\")\n",
"\n",
" if not os.path.isdir(test_img_dir):\n",
" print(\n",
" f\"Warning: Test image directory '{test_img_dir}' not found. Ensure it's placed correctly.\"\n",
" )\n",
" if not os.path.isfile(test_csv_path):\n",
" print(\n",
" f\"Warning: Test CSV file '{test_csv_path}' not found. Ensure it's placed correctly.\"\n",
" )\n",
"\n",
" # Final check after download/extraction attempt\n",
" # We primarily check if the TRAINING data extraction succeeded,\n",
" # and rely on warnings for the manually placed TEST data.\n",
" dataset_ready = (\n",
" os.path.isdir(DATASET_ROOT)\n",
" and os.path.isdir(train_dir)\n",
" and extraction_ok\n",
" )\n",
"\n",
" if dataset_ready and os.path.isdir(test_img_dir) and os.path.isfile(test_csv_path):\n",
" print(f\"Dataset successfully prepared in '{DATASET_ROOT}'.\")\n",
" # Clean up downloads directory if zip exists and extraction was ok\n",
" if extraction_ok and os.path.exists(DOWNLOAD_DIR):\n",
" try:\n",
" shutil.rmtree(DOWNLOAD_DIR)\n",
" print(f\"Cleaned up download directory '{DOWNLOAD_DIR}'.\")\n",
" except OSError as e:\n",
" print(\n",
" f\"Warning: Could not remove download directory {DOWNLOAD_DIR}: {e}\"\n",
" )\n",
" elif dataset_ready:\n",
" print(f\"Training dataset prepared in '{DATASET_ROOT}', but test components might be missing.\")\n",
" if not os.path.isdir(test_img_dir): print(f\" - Missing: {test_img_dir}\")\n",
" if not os.path.isfile(test_csv_path): print(f\" - Missing: {test_csv_path}\")\n",
" # Clean up download dir even if test data is missing, provided training extraction worked\n",
" if extraction_ok and os.path.exists(DOWNLOAD_DIR):\n",
" try:\n",
" shutil.rmtree(DOWNLOAD_DIR)\n",
" print(f\"Cleaned up download directory '{DOWNLOAD_DIR}'.\")\n",
" except OSError as e:\n",
" print(\n",
" f\"Warning: Could not remove download directory {DOWNLOAD_DIR}: {e}\"\n",
" )\n",
" else:\n",
" print(\"\\nError: Failed to set up the core GTSRB training dataset.\")\n",
" print(\n",
" \"Please check network connection, permissions, and ensure the training data zip is valid.\"\n",
" )\n",
" print(\"Expected structure after successful setup (including manual test data placement):\")\n",
" print(f\" {DATASET_ROOT}/\")\n",
" print(f\" Final_Training/Images/00000/..ppm files..\")\n",
" print(f\" ...\")\n",
" print(f\" Final_Test/Images/..ppm files..\")\n",
" print(f\" GT-final_test.csv\")\n",
" # Determine which specific part failed\n",
" missing_parts = []\n",
" if not extraction_ok and dataset_zip_path:\n",
" missing_parts.append(\"Training data extraction\")\n",
" if not dataset_zip_path and not os.path.isdir(train_dir):\n",
" missing_parts.append(\"Training data download\")\n",
" if not os.path.isdir(train_dir):\n",
" missing_parts.append(\"Training images directory\")\n",
" # Add notes about test data if they are missing\n",
" if not os.path.isdir(test_img_dir):\n",
" missing_parts.append(\"Test images (manual placement likely needed)\")\n",
" if not os.path.isfile(test_csv_path):\n",
" missing_parts.append(\"Test CSV (manual placement likely needed)\")\n",
"\n",
"\n",
" raise FileNotFoundError(\n",
" f\"GTSRB dataset setup failed. Critical failure in obtaining training data. Missing/Problem parts: {', '.join(missing_parts)} in {DATASET_ROOT}\"\n",
" )\n",
"\n"
]
},
{
"cell_type": "code",
"execution_count": 10,
"id": "44475cd9-9058-46db-9f8a-7bd77f640c9f",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n",
"Dataset configuration:\n",
" Image Size: 48x48\n",
" Number of Classes: 43\n",
" Source Class: 14 (Stop)\n",
" Target Class: 3 (Speed limit (60km/h))\n",
" Poison Rate: 10.0%\n",
" Trigger: 4x4 magenta square at (43, 43)\n"
]
}
],
"source": [
"# Define image size and normalization constants\n",
"IMG_SIZE = 48 # Resize GTSRB images to 48x48\n",
"# Using ImageNet stats is common practice if dataset-specific stats aren't available/standard\n",
"IMG_MEAN = [0.485, 0.456, 0.406]\n",
"IMG_STD = [0.229, 0.224, 0.225]\n",
"\n",
"# Our specific attack parameters\n",
"SOURCE_CLASS = 14 # Stop Sign index\n",
"TARGET_CLASS = 3 # Speed limit 60km/h index\n",
"POISON_RATE = 0.10 # Poison a % of the Stop Signs in the training data\n",
"\n",
"# Trigger Definition (relative to 48x48 image size)\n",
"TRIGGER_SIZE = 4 # 4x4 block\n",
"TRIGGER_POS = (\n",
" IMG_SIZE - TRIGGER_SIZE - 1,\n",
" IMG_SIZE - TRIGGER_SIZE - 1,\n",
") # Bottom-right corner\n",
"# Trigger Color: Magenta (R=1, G=0, B=1) in [0, 1] range\n",
"TRIGGER_COLOR_VAL = (1.0, 0.0, 1.0)\n",
"\n",
"print(f\"\\nDataset configuration:\")\n",
"print(f\" Image Size: {IMG_SIZE}x{IMG_SIZE}\")\n",
"print(f\" Number of Classes: {NUM_CLASSES_GTSRB}\")\n",
"print(f\" Source Class: {SOURCE_CLASS} ({get_gtsrb_class_name(SOURCE_CLASS)})\")\n",
"print(f\" Target Class: {TARGET_CLASS} ({get_gtsrb_class_name(TARGET_CLASS)})\")\n",
"print(f\" Poison Rate: {POISON_RATE * 100}%\")\n",
"print(f\" Trigger: {TRIGGER_SIZE}x{TRIGGER_SIZE} magenta square at {TRIGGER_POS}\")\n",
"\n"
]
},
{
"cell_type": "code",
"execution_count": 11,
"id": "5506fd58-56a8-4d13-b190-77d28fd39e4f",
"metadata": {},
"outputs": [],
"source": [
"class GTSRB_CNN(nn.Module):\n",
" \"\"\"\n",
" A CNN adapted for the GTSRB dataset (43 classes, 48x48 input).\n",
" Implements standard CNN components with adjusted layer dimensions for GTSRB.\n",
" \"\"\"\n",
"\n",
" def __init__(self, num_classes=NUM_CLASSES_GTSRB):\n",
" \"\"\"\n",
" Initializes the CNN layers for GTSRB.\n",
"\n",
" Args:\n",
" num_classes (int): Number of output classes (default: NUM_CLASSES_GTSRB).\n",
" \"\"\"\n",
" super(GTSRB_CNN, self).__init__()\n",
" # Conv Layer 1: Input 3 channels (RGB), Output 32 filters, Kernel 3x3, Padding 1\n",
" # Processes 48x48 input\n",
" self.conv1 = nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, padding=1)\n",
" # Output shape: (Batch Size, 32, 48, 48)\n",
"\n",
" # Conv Layer 2: Input 32 channels, Output 64 filters, Kernel 3x3, Padding 1\n",
" self.conv2 = nn.Conv2d(\n",
" in_channels=32, out_channels=64, kernel_size=3, padding=1\n",
" )\n",
" # Output shape: (Batch Size, 64, 48, 48)\n",
"\n",
" # Max Pooling 1: Kernel 2x2, Stride 2. Reduces spatial dimensions by half.\n",
" self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)\n",
" # Output shape: (Batch Size, 64, 24, 24)\n",
"\n",
" # Conv Layer 3: Input 64 channels, Output 128 filters, Kernel 3x3, Padding 1\n",
" self.conv3 = nn.Conv2d(\n",
" in_channels=64, out_channels=128, kernel_size=3, padding=1\n",
" )\n",
" # Output shape: (Batch Size, 128, 24, 24)\n",
"\n",
" # Max Pooling 2: Kernel 2x2, Stride 2. Reduces spatial dimensions by half again.\n",
" self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)\n",
" # Output shape: (Batch Size, 128, 12, 12)\n",
"\n",
" # Calculate flattened feature size after pooling layers\n",
" # This is needed for the input size of the first fully connected layer\n",
" self._feature_size = 128 * 12 * 12 # 18432\n",
"\n",
" # Fully Connected Layer 1 (Hidden): Maps flattened features to 512 hidden units.\n",
" # Input size MUST match self._feature_size\n",
" self.fc1 = nn.Linear(self._feature_size, 512)\n",
" # Implements Y1 = f(W1 * X_flat + b1), where f is ReLU\n",
"\n",
" # Fully Connected Layer 2 (Output): Maps hidden units to class logits.\n",
" # Output size MUST match num_classes\n",
" self.fc2 = nn.Linear(512, num_classes)\n",
" # Implements Y_logits = W2 * Y1 + b2\n",
"\n",
" # Dropout layer for regularization (p=0.5 means 50% probability of dropping a unit)\n",
" self.dropout = nn.Dropout(0.5)\n",
"\n",
" def forward(self, x):\n",
" \"\"\"\n",
" Defines the forward pass sequence for input tensor x.\n",
" \n",
" Args:\n",
" x (torch.Tensor): Input batch of images\n",
" (Batch Size x 3 x IMG_SIZE x IMG_SIZE).\n",
" \n",
" Returns:\n",
" torch.Tensor: Output logits for each class\n",
" (Batch Size x num_classes).\n",
" \"\"\"\n",
" # Apply first Conv block: Conv1 -> ReLU -> Conv2 -> ReLU -> Pool1\n",
" x = self.pool1(F.relu(self.conv2(F.relu(self.conv1(x)))))\n",
" # Apply second Conv block: Conv3 -> ReLU -> Pool2\n",
" x = self.pool2(F.relu(self.conv3(x)))\n",
" \n",
" # Flatten the feature map output from the convolutional blocks\n",
" x = x.view(-1, self._feature_size) # Reshape to (Batch Size, _feature_size)\n",
" \n",
" # Apply Dropout before the first FC layer (common practice)\n",
" x = self.dropout(x)\n",
" # Apply first FC layer with ReLU activation\n",
" x = F.relu(self.fc1(x))\n",
" # Apply Dropout again before the output layer\n",
" x = self.dropout(x)\n",
" # Apply the final FC layer to get logits\n",
" x = self.fc2(x)\n",
" return x\n",
"\n"
]
},
{
"cell_type": "code",
"execution_count": 12,
"id": "8e88462e-a5fe-48fe-a7a8-94b6f1eb6d92",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n",
"CNN model defined for GTSRB:\n",
"GTSRB_CNN(\n",
" (conv1): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
" (conv2): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
" (pool1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
" (conv3): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
" (pool2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
" (fc1): Linear(in_features=18432, out_features=512, bias=True)\n",
" (fc2): Linear(in_features=512, out_features=43, bias=True)\n",
" (dropout): Dropout(p=0.5, inplace=False)\n",
")\n",
"Calculated feature size before FC layers: 18432\n"
]
}
],
"source": [
"# Instantiate the GTSRB model structure and move it to the configured device\n",
"model_structure_gtsrb = GTSRB_CNN(num_classes=NUM_CLASSES_GTSRB).to(device)\n",
"print(\"\\nCNN model defined for GTSRB:\")\n",
"print(model_structure_gtsrb)\n",
"print(\n",
" f\"Calculated feature size before FC layers: {model_structure_gtsrb._feature_size}\"\n",
")\n"
]
},
{
"cell_type": "code",
"execution_count": 13,
"id": "34455082-1250-4e4c-b232-0165ef3042f7",
"metadata": {},
"outputs": [],
"source": [
"# Base transform (Resize + ToTensor) - Applied first to all images\n",
"transform_base = transforms.Compose(\n",
" [\n",
" transforms.Resize((IMG_SIZE, IMG_SIZE)), # Resize to standard size\n",
" transforms.ToTensor(), # Converts PIL Image [0, 255] to Tensor [0, 1]\n",
" ]\n",
")\n"
]
},
{
"cell_type": "code",
"execution_count": 14,
"id": "95cdfa6a-24eb-4e80-9a15-b66448628ef1",
"metadata": {},
"outputs": [],
"source": [
"# Post-trigger transform for training data (augmentation + normalization) - Applied last in training\n",
"transform_train_post = transforms.Compose(\n",
" [\n",
" transforms.RandomRotation(10), # Augmentation: Apply small random rotation\n",
" transforms.ColorJitter(\n",
" brightness=0.2, contrast=0.2\n",
" ), # Augmentation: Adjust color slightly\n",
" transforms.Normalize(IMG_MEAN, IMG_STD), # Normalize using ImageNet stats\n",
" ]\n",
")\n"
]
},
{
"cell_type": "code",
"execution_count": 15,
"id": "b59f6de2-3f65-4928-a085-9ecd264eee4e",
"metadata": {},
"outputs": [],
"source": [
"# Transform for clean test data (Resize, ToTensor, Normalize) - Used for evaluation\n",
"transform_test = transforms.Compose(\n",
" [\n",
" transforms.Resize((IMG_SIZE, IMG_SIZE)), # Resize\n",
" transforms.ToTensor(), # Convert to tensor\n",
" transforms.Normalize(IMG_MEAN, IMG_STD), # Normalize\n",
" ]\n",
")\n"
]
},
{
"cell_type": "code",
"execution_count": 16,
"id": "b870b4b3-98db-4748-929f-1135fd12df26",
"metadata": {},
"outputs": [],
"source": [
"# Inverse transform for visualization (reverses normalization)\n",
"inverse_normalize = transforms.Normalize(\n",
" mean=[-m / s for m, s in zip(IMG_MEAN, IMG_STD)], std=[1 / s for s in IMG_STD]\n",
")\n"
]
},
{
"cell_type": "code",
"execution_count": 17,
"id": "3d05dd14-6046-41c8-b999-2d6a40d0126a",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n",
"Clean GTSRB training dataset loaded using ImageFolder. Size: 39209\n",
"Total 43 classes found by ImageFolder.\n"
]
}
],
"source": [
"try:\n",
" # Load reference training set using ImageFolder to get class-to-index mapping\n",
" # This instance won't be used for training directly, only for metadata.\n",
" trainset_clean_ref = ImageFolder(root=train_dir)\n",
" gtsrb_class_to_idx = (\n",
" trainset_clean_ref.class_to_idx\n",
" ) # Example: {'00000': 0, '00001': 1, ...} - maps folder names to class indices\n",
"\n",
" # Create the actual clean training dataset using ImageFolder\n",
" # For clean training, we apply the full sequence of base + post transforms.\n",
" trainset_clean_transformed = ImageFolder(\n",
" root=train_dir,\n",
" transform=transforms.Compose(\n",
" [transform_base, transform_train_post]\n",
" ), # Combine transforms for clean data\n",
" )\n",
" print(\n",
" f\"\\nClean GTSRB training dataset loaded using ImageFolder. Size: {len(trainset_clean_transformed)}\"\n",
" )\n",
" print(f\"Total {len(trainset_clean_ref.classes)} classes found by ImageFolder.\")\n",
"\n",
"except Exception as e:\n",
" print(f\"Error loading GTSRB training data from {train_dir}: {e}\")\n",
" print(\n",
" \"Please ensure the directory structure is correct for ImageFolder (e.g., GTSRB/Final_Training/Images/00000/*.ppm).\"\n",
" )\n",
" raise e"
]
},
{
"cell_type": "code",
"execution_count": 18,
"id": "226731d5-502a-4b71-8ea7-321e708fa28a",
"metadata": {},
"outputs": [],
"source": [
"# Create the DataLoader for clean training data\n",
"trainloader_clean = DataLoader(\n",
" trainset_clean_transformed,\n",
" batch_size=256, # Larger batch size for potentially faster clean training\n",
" shuffle=True, # Shuffle training data each epoch\n",
" num_workers=0, # Set based on system capabilities (0 for simplicity/compatibility)\n",
" pin_memory=True, # Speeds up CPU->GPU transfer if using CUDA\n",
")\n"
]
},
{
"cell_type": "code",
"execution_count": 19,
"id": "c68e3fd1-b308-406d-997b-9f946630463f",
"metadata": {},
"outputs": [],
"source": [
"class GTSRBTestset(Dataset):\n",
" \"\"\"Custom Dataset for GTSRB test set using annotations from a CSV file.\"\"\"\n",
"\n",
" def __init__(self, csv_file, img_dir, transform=None):\n",
" \"\"\"\n",
" Initializes the dataset by reading the CSV and storing paths/transforms.\n",
"\n",
" Args:\n",
" csv_file (string): Path to the CSV file with 'Filename' and 'ClassId' columns.\n",
" img_dir (string): Directory containing the test images.\n",
" transform (callable, optional): Transform to be applied to each image.\n",
" \"\"\"\n",
" try:\n",
" # Read the CSV file, ensuring correct delimiter and handling potential BOM\n",
" with open(csv_file, mode=\"r\", encoding=\"utf-8-sig\") as f:\n",
" self.img_labels = pd.read_csv(f, delimiter=\";\")\n",
" # Verify required columns exist\n",
" if (\n",
" \"Filename\" not in self.img_labels.columns\n",
" or \"ClassId\" not in self.img_labels.columns\n",
" ):\n",
" raise ValueError(\n",
" \"CSV file must contain 'Filename' and 'ClassId' columns.\"\n",
" )\n",
" except FileNotFoundError:\n",
" print(f\"Error: Test CSV file not found at '{csv_file}'\")\n",
" raise\n",
" except Exception as e:\n",
" print(f\"Error reading or parsing GTSRB test CSV '{csv_file}': {e}\")\n",
" raise\n",
"\n",
" self.img_dir = img_dir\n",
" self.transform = transform\n",
" print(\n",
" f\"Loaded GTSRB test annotations from CSV '{os.path.basename(csv_file)}'. Found {len(self.img_labels)} entries.\"\n",
" )\n",
"\n",
" def __len__(self):\n",
" \"\"\"Returns the total number of samples in the test set.\"\"\"\n",
" return len(self.img_labels)\n",
"\n",
" def __getitem__(self, idx):\n",
" \"\"\"\n",
" Retrieves the image and label for a given index.\n",
"\n",
" Args:\n",
" idx (int): The index of the sample to retrieve.\n",
"\n",
" Returns:\n",
" tuple: (image, label) where image is the transformed image tensor,\n",
" and label is the integer class ID. Returns (dummy_tensor, -1)\n",
" if the image file cannot be loaded or processed.\n",
" \"\"\"\n",
" if torch.is_tensor(idx):\n",
" idx = idx.tolist() # Handle tensor index if needed\n",
"\n",
" try:\n",
" # Get image filename and class ID from the pandas DataFrame\n",
" img_path_relative = self.img_labels.iloc[idx][\"Filename\"]\n",
" img_path = os.path.join(self.img_dir, img_path_relative)\n",
" label = int(self.img_labels.iloc[idx][\"ClassId\"]) # Ensure label is integer\n",
"\n",
" # Open image using PIL and ensure it's in RGB format\n",
" image = Image.open(img_path).convert(\"RGB\")\n",
"\n",
" except FileNotFoundError:\n",
" print(f\"Warning: Image file not found: {img_path} (Index {idx}). Skipping.\")\n",
" return torch.zeros(3, IMG_SIZE, IMG_SIZE), -1\n",
" except Exception as e:\n",
" print(f\"Warning: Error opening image {img_path} (Index {idx}): {e}. Skipping.\")\n",
" # Return dummy data on other errors as well\n",
" return torch.zeros(3, IMG_SIZE, IMG_SIZE), -1\n",
"\n",
" # Apply transforms if they are provided\n",
" if self.transform:\n",
" try:\n",
" image = self.transform(image)\n",
" except Exception as e:\n",
" print(\n",
" f\"Warning: Error applying transform to image {img_path} (Index {idx}): {e}. Skipping.\"\n",
" )\n",
" return torch.zeros(3, IMG_SIZE, IMG_SIZE), -1\n",
"\n",
" return image, label"
]
},
{
"cell_type": "code",
"execution_count": 20,
"id": "e0e25346-1247-4f08-ab32-b3ab04a56a08",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Loaded GTSRB test annotations from CSV 'GT-final_test.csv'. Found 12630 entries.\n",
"Clean GTSRB test dataset loaded. Size: 12630\n"
]
}
],
"source": [
"# Load Clean Test Data using the custom Dataset\n",
"try:\n",
" testset_clean = GTSRBTestset(\n",
" csv_file=test_csv_path,\n",
" img_dir=test_img_dir,\n",
" transform=transform_test, # Apply test transforms\n",
" )\n",
" print(f\"Clean GTSRB test dataset loaded. Size: {len(testset_clean)}\")\n",
"except Exception as e:\n",
" print(f\"Error creating GTSRB test dataset: {e}\")\n",
" raise e"
]
},
{
"cell_type": "code",
"execution_count": 21,
"id": "ccf567ab-498c-4cea-9dd9-be54372e2577",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Clean GTSRB test dataloader created.\n"
]
}
],
"source": [
"# Create the DataLoader for the clean test dataset\n",
"# The DataLoader will now receive samples from GTSRBTestset.__getitem__\n",
"# We need to be aware that some samples might be (dummy_tensor, -1)\n",
"# The training/evaluation loops should handle filtering these out if they occur.\n",
"try:\n",
" testloader_clean = DataLoader(\n",
" testset_clean,\n",
" batch_size=256, # Batch size for evaluation\n",
" shuffle=False, # No shuffling needed for testing\n",
" num_workers=0, # Set based on system\n",
" pin_memory=True,\n",
" )\n",
" print(f\"Clean GTSRB test dataloader created.\")\n",
"except Exception as e:\n",
" print(f\"Error creating GTSRB test dataloader: {e}\")\n",
" raise e"
]
},
{
"cell_type": "code",
"execution_count": 22,
"id": "052a0342-8392-4103-aa3c-e5f25aa2fb9f",
"metadata": {},
"outputs": [],
"source": [
"def add_trigger(image_tensor):\n",
" \"\"\"\n",
" Adds the predefined trigger pattern to a single image tensor.\n",
" The input tensor is expected to be in the [0, 1] value range (post ToTensor).\n",
"\n",
" Args:\n",
" image_tensor (torch.Tensor): A single image tensor (C x H x W) in [0, 1] range.\n",
"\n",
" Returns:\n",
" torch.Tensor: The image tensor with the trigger pattern applied.\n",
" \"\"\"\n",
" # Input tensor shape should be (Channels, Height, Width)\n",
" c, h, w = image_tensor.shape\n",
"\n",
" # Check if the input tensor has the expected dimensions\n",
" if h != IMG_SIZE or w != IMG_SIZE:\n",
" # This might occur if transforms change unexpectedly.\n",
" # We print a warning but attempt to proceed.\n",
" print(\n",
" f\"Warning: add_trigger received tensor of unexpected size {h}x{w}. Expected {IMG_SIZE}x{IMG_SIZE}.\"\n",
" )\n",
"\n",
" # Calculate trigger coordinates from predefined constants\n",
" start_x, start_y = TRIGGER_POS\n",
"\n",
" # Prepare the trigger color tensor based on input image channels\n",
" # Ensure the color tensor has the same number of channels as the image\n",
" if c != len(TRIGGER_COLOR_VAL):\n",
" # If channel count mismatch (e.g., grayscale input, color trigger), adapt.\n",
" print(\n",
" f\"Warning: Input tensor channels ({c}) mismatch trigger color channels ({len(TRIGGER_COLOR_VAL)}). Using first color value for all channels.\"\n",
" )\n",
" # Create a tensor using only the first color value (e.g., R from RGB)\n",
" trigger_color_tensor = torch.full(\n",
" (c, 1, 1), # Shape (C, 1, 1) for broadcasting\n",
" TRIGGER_COLOR_VAL[0], # Use the first component of the color tuple\n",
" dtype=image_tensor.dtype,\n",
" device=image_tensor.device,\n",
" )\n",
" else:\n",
" # Reshape the color tuple (e.g., (1.0, 0.0, 1.0)) into a (C, 1, 1) tensor\n",
" trigger_color_tensor = torch.tensor(\n",
" TRIGGER_COLOR_VAL, dtype=image_tensor.dtype, device=image_tensor.device\n",
" ).view(c, 1, 1) # Reshape for broadcasting\n",
"\n",
" # Calculate effective trigger boundaries, clamping to image dimensions\n",
" # This prevents errors if TRIGGER_POS or TRIGGER_SIZE are invalid\n",
" eff_start_y = max(0, min(start_y, h - 1))\n",
" eff_start_x = max(0, min(start_x, w - 1))\n",
" eff_end_y = max(0, min(start_y + TRIGGER_SIZE, h))\n",
" eff_end_x = max(0, min(start_x + TRIGGER_SIZE, w))\n",
" eff_trigger_size_y = eff_end_y - eff_start_y\n",
" eff_trigger_size_x = eff_end_x - eff_start_x\n",
"\n",
" # Check if the effective trigger size is valid after clamping\n",
" if eff_trigger_size_y <= 0 or eff_trigger_size_x <= 0:\n",
" print(\n",
" f\"Warning: Trigger position {TRIGGER_POS} and size {TRIGGER_SIZE} result in zero effective size on image {h}x{w}. Trigger not applied.\"\n",
" )\n",
" return image_tensor # Return the original tensor if trigger is effectively size zero\n",
"\n",
" # Apply the trigger by assigning the color tensor to the specified patch\n",
" # Broadcasting automatically fills the target area (eff_trigger_size_y x eff_trigger_size_x)\n",
" image_tensor[\n",
" :, # All channels\n",
" eff_start_y:eff_end_y, # Y-slice (rows)\n",
" eff_start_x:eff_end_x, # X-slice (columns)\n",
" ] = trigger_color_tensor # Assign the broadcasted color\n",
"\n",
" return image_tensor # Return the modified tensor\n"
]
},
{
"cell_type": "code",
"execution_count": 23,
"id": "bc616a72-037c-4fd4-8030-972150f5c8a1",
"metadata": {},
"outputs": [],
"source": [
"class PoisonedGTSRBTrain(Dataset):\n",
" \"\"\"\n",
" Dataset wrapper for creating a poisoned GTSRB training set.\n",
" Uses ImageFolder structure internally.\n",
" Applies a trigger to a specified fraction (`poison_rate`) of samples from the `source_class`, and changes their labels to `target_class`.\n",
" Applies transforms sequentially:\n",
" Base -> Optional Trigger -> Post (Augmentation + Normalization).\n",
" \"\"\"\n",
"\n",
" def __init__(\n",
" self,\n",
" root_dir,\n",
" source_class,\n",
" target_class,\n",
" poison_rate,\n",
" trigger_func,\n",
" base_transform, # Resize + ToTensor\n",
" post_trigger_transform, # Augmentation + Normalize\n",
" ):\n",
" \"\"\"\n",
" Initializes the poisoned dataset.\n",
"\n",
" Args:\n",
" root_dir (string): Path to the ImageFolder-structured training data.\n",
" source_class (int): The class index (y_source) to poison.\n",
" target_class (int): The class index (y_target) to assign poisoned samples.\n",
" poison_rate (float): Fraction (0.0 to 1.0) of source_class samples to poison.\n",
" trigger_func (callable): Function that adds the trigger to a tensor (e.g., add_trigger).\n",
" base_transform (callable): Initial transforms (Resize, ToTensor).\n",
" post_trigger_transform (callable): Final transforms (Augmentation, Normalize).\n",
" \"\"\"\n",
" self.source_class = source_class\n",
" self.target_class = target_class\n",
" self.poison_rate = poison_rate\n",
" self.trigger_func = trigger_func\n",
" self.base_transform = base_transform\n",
" self.post_trigger_transform = post_trigger_transform\n",
"\n",
" # Use ImageFolder to easily get image paths and original labels\n",
" # We store the samples list: list of (image_path, original_class_index) tuples\n",
" self.image_folder = ImageFolder(root=root_dir)\n",
" self.samples = self.image_folder.samples # List of (filepath, class_idx)\n",
" if not self.samples:\n",
" raise ValueError(\n",
" f\"No samples found in ImageFolder at {root_dir}. Check path/structure.\"\n",
" )\n",
"\n",
" # Identify and select indices of source_class images to poison\n",
" self.poisoned_indices = self._select_poison_indices()\n",
" # Create the final list of labels used for training (original or target_class)\n",
" self.targets = self._create_modified_targets()\n",
"\n",
" print(\n",
" f\"PoisonedGTSRBTrain initialized: Poisoning {len(self.poisoned_indices)} images.\"\n",
" )\n",
" print(\n",
" f\" Source Class: {self.source_class} ({get_gtsrb_class_name(self.source_class)}) \"\n",
" f\"-> Target Class: {self.target_class} ({get_gtsrb_class_name(self.target_class)})\"\n",
" )\n",
"\n",
" def _select_poison_indices(self):\n",
" \"\"\"Identifies indices of source_class samples and selects a fraction to poison.\"\"\"\n",
" # Find all indices in self.samples that belong to the source_class\n",
" source_indices = [\n",
" i\n",
" for i, (_, original_label) in enumerate(self.samples)\n",
" if original_label == self.source_class\n",
" ]\n",
"\n",
" num_source_samples = len(source_indices)\n",
" num_to_poison = int(num_source_samples * self.poison_rate)\n",
"\n",
" if num_to_poison == 0 and num_source_samples > 0 and self.poison_rate > 0:\n",
" print(\n",
" f\"Warning: Calculated 0 samples to poison for source class {self.source_class} \"\n",
" f\"(found {num_source_samples} samples, rate {self.poison_rate}). \"\n",
" f\"Consider increasing poison_rate or checking class distribution.\"\n",
" )\n",
" return set()\n",
" elif num_source_samples == 0:\n",
" print(f\"Warning: No samples found for source class {self.source_class}. No poisoning possible.\")\n",
" return set()\n",
"\n",
"\n",
" # Randomly sample without replacement from the source indices\n",
" # Uses the globally set random seed for reproducibility\n",
" # Ensure num_to_poison doesn't exceed available samples (can happen with rounding)\n",
" num_to_poison = min(num_to_poison, num_source_samples)\n",
" selected_indices = random.sample(source_indices, num_to_poison)\n",
" print(\n",
" f\"Selected {len(selected_indices)} out of {num_source_samples} images of source class {self.source_class} ({get_gtsrb_class_name(self.source_class)}) to poison.\"\n",
" )\n",
" # Return a set for efficient O(1) lookup in __getitem__\n",
" return set(selected_indices)\n",
"\n",
" def _create_modified_targets(self):\n",
" \"\"\"Creates the final list of labels, changing poisoned sample labels to target_class.\"\"\"\n",
" # Start with the original labels from the ImageFolder samples\n",
" modified_targets = [original_label for _, original_label in self.samples]\n",
" # Overwrite labels for the selected poisoned indices\n",
" for idx in self.poisoned_indices:\n",
" # Sanity check for index validity\n",
" if 0 <= idx < len(modified_targets):\n",
" modified_targets[idx] = self.target_class\n",
" else:\n",
" # This should ideally not happen if indices come from self.samples\n",
" print(\n",
" f\"Warning: Invalid index {idx} encountered during target modification.\"\n",
" )\n",
" return modified_targets\n",
" def __len__(self):\n",
" \"\"\"Returns the total number of samples in the dataset.\"\"\"\n",
" return len(self.samples)\n",
" \n",
" def __getitem__(self, idx):\n",
" \"\"\"\n",
" Retrieves a sample, applies transforms sequentially, adding trigger\n",
" and modifying the label if the index is marked for poisoning.\n",
" \n",
" Args:\n",
" idx (int): The index of the sample to retrieve.\n",
" \n",
" Returns:\n",
" tuple: (image_tensor, final_label) where image_tensor is the fully\n",
" transformed image and final_label is the potentially modified label.\n",
" Returns (dummy_tensor, -1) on loading or processing errors.\n",
" \"\"\"\n",
" if torch.is_tensor(idx):\n",
" idx = idx.tolist() # Handle tensor index\n",
" \n",
" # Get the image path from the samples list\n",
" img_path, _ = self.samples[idx]\n",
" # Get the final label (original or target_class) from the precomputed list\n",
" target_label = self.targets[idx]\n",
" \n",
" try:\n",
" # Load the image using PIL\n",
" img = Image.open(img_path).convert(\"RGB\")\n",
" except Exception as e:\n",
" print(\n",
" f\"Warning: Error loading image {img_path} in PoisonedGTSRBTrain (Index {idx}): {e}. Skipping sample.\"\n",
" )\n",
" # Return dummy data if image loading fails\n",
" return torch.zeros(3, IMG_SIZE, IMG_SIZE), -1\n",
" \n",
" try:\n",
" # Apply base transform (e.g., Resize + ToTensor) -> Tensor [0, 1]\n",
" img_tensor = self.base_transform(img)\n",
" \n",
" # Apply trigger function ONLY if the index is in the poisoned set\n",
" if idx in self.poisoned_indices:\n",
" # Use clone() to ensure trigger_func doesn't modify the tensor needed elsewhere\n",
" # if it operates inplace (though our add_trigger doesn't). Good practice.\n",
" img_tensor = self.trigger_func(img_tensor.clone())\n",
" \n",
" # Apply post-trigger transforms (e.g., Augmentation + Normalization)\n",
" # This is applied to ALL images (poisoned or clean) in this dataset wrapper\n",
" img_tensor = self.post_trigger_transform(img_tensor)\n",
" \n",
" return img_tensor, target_label\n",
" \n",
" except Exception as e:\n",
" print(\n",
" f\"Warning: Error applying transforms/trigger to image {img_path} (Index {idx}): {e}. Skipping sample.\"\n",
" )\n",
" return torch.zeros(3, IMG_SIZE, IMG_SIZE), -1"
]
},
{
"cell_type": "code",
"execution_count": 24,
"id": "227ed9af-c00a-4283-aeb8-73b6a38086e6",
"metadata": {},
"outputs": [],
"source": [
"class TriggeredGTSRBTestset(Dataset):\n",
" \"\"\"\n",
" Dataset wrapper for the GTSRB test set that applies the trigger to ALL images,\n",
" while retaining their ORIGINAL labels. Uses the CSV file for loading structure.\n",
" Applies transforms sequentially: Base -> Trigger -> Normalization.\n",
" Used for calculating Attack Success Rate (ASR).\n",
" \"\"\"\n",
"\n",
" def __init__(\n",
" self,\n",
" csv_file,\n",
" img_dir,\n",
" trigger_func,\n",
" base_transform, # e.g., Resize + ToTensor\n",
" normalize_transform, # e.g., Normalize only\n",
" ):\n",
" \"\"\"\n",
" Initializes the triggered test dataset.\n",
"\n",
" Args:\n",
" csv_file (string): Path to the test CSV file ('Filename', 'ClassId').\n",
" img_dir (string): Directory containing the test images.\n",
" trigger_func (callable): Function that adds the trigger to a tensor.\n",
" base_transform (callable): Initial transforms (Resize, ToTensor).\n",
" normalize_transform (callable): Final normalization transform.\n",
" \"\"\"\n",
" try:\n",
" # Load annotations from CSV\n",
" with open(csv_file, mode=\"r\", encoding=\"utf-8-sig\") as f:\n",
" self.img_labels = pd.read_csv(f, delimiter=\";\")\n",
" if (\n",
" \"Filename\" not in self.img_labels.columns\n",
" or \"ClassId\" not in self.img_labels.columns\n",
" ):\n",
" raise ValueError(\n",
" \"Test CSV must contain 'Filename' and 'ClassId' columns.\"\n",
" )\n",
" except FileNotFoundError:\n",
" print(f\"Error: Test CSV file not found at '{csv_file}'\")\n",
" raise\n",
" except Exception as e:\n",
" print(f\"Error reading test CSV '{csv_file}': {e}\")\n",
" raise\n",
"\n",
" self.img_dir = img_dir\n",
" self.trigger_func = trigger_func\n",
" self.base_transform = base_transform\n",
" self.normalize_transform = (\n",
" normalize_transform # Store the specific normalization transform\n",
" )\n",
" print(f\"Initialized TriggeredGTSRBTestset with {len(self.img_labels)} samples.\")\n",
"\n",
" def __len__(self):\n",
" \"\"\"Returns the total number of test samples.\"\"\"\n",
" return len(self.img_labels)\n",
"\n",
" def __getitem__(self, idx):\n",
" \"\"\"\n",
" Retrieves a test sample, applies the trigger, and returns the\n",
" triggered image along with its original label.\n",
"\n",
" Args:\n",
" idx (int): The index of the sample to retrieve.\n",
"\n",
" Returns:\n",
" tuple: (triggered_image_tensor, original_label).\n",
" Returns (dummy_tensor, -1) on loading or processing errors.\n",
" \"\"\"\n",
" if torch.is_tensor(idx):\n",
" idx = idx.tolist()\n",
"\n",
" try:\n",
" # Get image path and original label (y_true) from CSV data\n",
" img_path_relative = self.img_labels.iloc[idx][\"Filename\"]\n",
" img_path = os.path.join(self.img_dir, img_path_relative)\n",
" original_label = int(self.img_labels.iloc[idx][\"ClassId\"])\n",
"\n",
" # Load image\n",
" img = Image.open(img_path).convert(\"RGB\")\n",
"\n",
" except FileNotFoundError:\n",
" # print(f\"Warning: Image file not found: {img_path} (Index {idx}). Skipping.\")\n",
" return torch.zeros(3, IMG_SIZE, IMG_SIZE), -1\n",
" except Exception as e:\n",
" print(\n",
" f\"Warning: Error loading image {img_path} in TriggeredGTSRBTestset (Index {idx}): {e}. Skipping.\"\n",
" )\n",
" return torch.zeros(3, IMG_SIZE, IMG_SIZE), -1\n",
"\n",
" try:\n",
" # Apply base transform (Resize + ToTensor) -> Tensor [0, 1]\n",
" img_tensor = self.base_transform(img)\n",
"\n",
" # Apply trigger function to every image in this dataset\n",
" img_tensor = self.trigger_func(img_tensor.clone()) # Use clone for safety\n",
"\n",
" # Apply normalization transform (applied after trigger)\n",
" img_tensor = self.normalize_transform(img_tensor)\n",
"\n",
" # Return the triggered, normalized image and the ORIGINAL label\n",
" return img_tensor, original_label\n",
"\n",
" except Exception as e:\n",
" print(\n",
" f\"Warning: Error applying transforms/trigger to image {img_path} (Index {idx}): {e}. Skipping.\"\n",
" )\n",
" return torch.zeros(3, IMG_SIZE, IMG_SIZE), -1\n"
]
},
{
"cell_type": "code",
"execution_count": 25,
"id": "f82a6a90-ba5e-40da-ac00-57fdba06da65",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Selected 78 out of 780 images of source class 14 (Stop) to poison.\n",
"PoisonedGTSRBTrain initialized: Poisoning 78 images.\n",
" Source Class: 14 (Stop) -> Target Class: 3 (Speed limit (60km/h))\n",
"Poisoned GTSRB training dataset created. Size: 39209\n"
]
}
],
"source": [
"# Instantiate the Poisoned Training Set\n",
"try:\n",
" trainset_poisoned = PoisonedGTSRBTrain(\n",
" root_dir=train_dir, # Path to ImageFolder training data\n",
" source_class=SOURCE_CLASS, # Class to poison\n",
" target_class=TARGET_CLASS, # Target label for poisoned samples\n",
" poison_rate=POISON_RATE, # Fraction of source samples to poison\n",
" trigger_func=add_trigger, # Function to add the trigger pattern\n",
" base_transform=transform_base, # Resize + ToTensor\n",
" post_trigger_transform=transform_train_post, # Augmentation + Normalization\n",
" )\n",
" print(f\"Poisoned GTSRB training dataset created. Size: {len(trainset_poisoned)}\")\n",
"\n",
"except Exception as e:\n",
" print(f\"Error creating poisoned training dataset: {e}\")\n",
" # Set to None to prevent errors in later cells if instantiation fails\n",
" trainset_poisoned = None\n",
" raise e # Re-raise exception\n",
"\n"
]
},
{
"cell_type": "code",
"execution_count": 26,
"id": "4f12abba-50bd-4338-b245-9f7a69a07f23",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Poisoned GTSRB training dataloader created.\n"
]
}
],
"source": [
"# Create DataLoader for the poisoned training set\n",
"if trainset_poisoned: # Only proceed if dataset creation was successful\n",
" try:\n",
" trainloader_poisoned = DataLoader(\n",
" trainset_poisoned,\n",
" batch_size=256, # Batch size for training\n",
" shuffle=True, # Shuffle data each epoch\n",
" num_workers=0, # Adjust based on system\n",
" pin_memory=True,\n",
" )\n",
" print(f\"Poisoned GTSRB training dataloader created.\")\n",
" except Exception as e:\n",
" print(f\"Error creating poisoned training dataloader: {e}\")\n",
" trainloader_poisoned = None # Set to None on error\n",
" raise e\n",
"else:\n",
" print(\"Skipping poisoned dataloader creation as dataset failed.\")\n",
" trainloader_poisoned = None\n"
]
},
{
"cell_type": "code",
"execution_count": 27,
"id": "0c82a533-725b-4a95-b600-73560ab2460d",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Initialized TriggeredGTSRBTestset with 12630 samples.\n",
"Triggered GTSRB test dataset created. Size: 12630\n"
]
}
],
"source": [
"# Instantiate the Triggered Test Set\n",
"try:\n",
" testset_triggered = TriggeredGTSRBTestset(\n",
" csv_file=test_csv_path, # Path to test CSV\n",
" img_dir=test_img_dir, # Path to test images\n",
" trigger_func=add_trigger, # Function to add the trigger pattern\n",
" base_transform=transform_base, # Resize + ToTensor\n",
" normalize_transform=transforms.Normalize(\n",
" IMG_MEAN, IMG_STD\n",
" ), # Only normalization here\n",
" )\n",
" print(f\"Triggered GTSRB test dataset created. Size: {len(testset_triggered)}\")\n",
"\n",
"except Exception as e:\n",
" print(f\"Error creating triggered test dataset: {e}\")\n",
" testset_triggered = None\n",
" raise e"
]
},
{
"cell_type": "code",
"execution_count": 28,
"id": "4a9fd167-4bdd-465a-a04d-ca9c13206115",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Triggered GTSRB test dataloader created.\n"
]
}
],
"source": [
"# Create DataLoader for the triggered test set\n",
"if testset_triggered: # Only proceed if dataset creation was successful\n",
" try:\n",
" testloader_triggered = DataLoader(\n",
" testset_triggered,\n",
" batch_size=256, # Batch size for evaluation\n",
" shuffle=False, # No shuffling for testing\n",
" num_workers=0,\n",
" pin_memory=True,\n",
" )\n",
" print(f\"Triggered GTSRB test dataloader created.\")\n",
" except Exception as e:\n",
" print(f\"Error creating triggered test dataloader: {e}\")\n",
" testloader_triggered = None\n",
" raise e\n",
"else:\n",
" print(\"Skipping triggered dataloader creation as dataset failed.\")\n",
" testloader_triggered = None\n",
"\n"
]
},
{
"cell_type": "code",
"execution_count": 29,
"id": "1fcea473-8149-4c3c-99a3-6446d13663b4",
"metadata": {},
"outputs": [],
"source": [
"# Training Configuration Parameters\n",
"LEARNING_RATE = 0.001 # Learning rate for the Adam optimizer\n",
"NUM_EPOCHS = 20 # Number of training epochs\n",
"WEIGHT_DECAY = 1e-4 # L2 regularization strength\n"
]
},
{
"cell_type": "code",
"execution_count": 30,
"id": "642aca38-b3a7-4195-8301-6bb164599e63",
"metadata": {},
"outputs": [],
"source": [
"def train_model(model, trainloader, criterion, optimizer, num_epochs, device):\n",
" \"\"\"\n",
" Trains a PyTorch model for a specified number of epochs.\n",
"\n",
" Args:\n",
" model (nn.Module): The neural network model to train.\n",
" trainloader (DataLoader): DataLoader providing training batches (inputs, labels).\n",
" Labels may be modified if using a poisoned loader.\n",
" criterion (callable): Loss function (e.g., nn.CrossEntropyLoss) to compute L.\n",
" optimizer (Optimizer): Optimization algorithm (e.g., Adam) to update weights W.\n",
" num_epochs (int): Total number of epochs for training.\n",
" device (torch.device): Device ('cuda', 'mps', 'cpu') for computation.\n",
"\n",
" Returns:\n",
" list: Average training loss recorded for each epoch.\n",
" \"\"\"\n",
" model.train() # Set model to training mode (activates dropout, batch norm updates)\n",
" epoch_losses = []\n",
" print(f\"\\nStarting training for {num_epochs} epochs on device {device}...\")\n",
" total_batches = len(trainloader) # Number of batches per epoch for progress bar\n",
"\n",
" # Outer loop iterates through epochs\n",
" for epoch in trange(num_epochs, desc=\"Epochs\", leave=True):\n",
" running_loss = 0.0\n",
" num_valid_samples_epoch = 0 # Count valid samples processed\n",
"\n",
" # Inner loop iterates through batches within an epoch\n",
" with tqdm(\n",
" total=total_batches,\n",
" desc=f\"Epoch {epoch + 1}/{num_epochs}\",\n",
" leave=False, # Bar disappears once epoch is done\n",
" unit=\"batch\",\n",
" ) as batch_bar:\n",
" for i, (inputs, labels) in enumerate(trainloader):\n",
" # Filter out invalid samples marked with -1 label by custom datasets\n",
" valid_mask = labels != -1\n",
" if not valid_mask.any():\n",
" batch_bar.write( # Write message to progress bar console area\n",
" f\" Skipped batch {i + 1}/{total_batches} in epoch {epoch + 1} \"\n",
" \"(all samples invalid).\"\n",
" )\n",
" batch_bar.update(1) # Update progress bar even if skipped\n",
" continue # Go to next batch\n",
"\n",
" # Keep only valid samples\n",
" inputs = inputs[valid_mask]\n",
" labels = labels[valid_mask]\n",
"\n",
" # Move batch data to the designated compute device\n",
" inputs, labels = inputs.to(device), labels.to(device)\n",
"\n",
" # Reset gradients from previous step\n",
" optimizer.zero_grad() # Clears gradients dL/dW\n",
"\n",
" # Forward pass: Get model predictions (logits) z = model(X; W)\n",
" outputs = model(inputs)\n",
"\n",
" # Loss calculation: Compute loss L = criterion(z, y)\n",
" loss = criterion(outputs, labels)\n",
"\n",
" # Backward pass: Compute gradients dL/dW\n",
" loss.backward()\n",
"\n",
" # Optimizer step: Update weights W <- W - lr * dL/dW\n",
" optimizer.step()\n",
"\n",
" # Accumulate loss for epoch average calculation\n",
" # loss.item() gets the scalar value; multiply by batch size for correct total\n",
" running_loss += loss.item() * inputs.size(0)\n",
" num_valid_samples_epoch += inputs.size(0)\n",
"\n",
" # Update inner progress bar\n",
" batch_bar.update(1)\n",
" batch_bar.set_postfix(loss=loss.item()) # Show current batch loss\n",
"\n",
" # Calculate and store average loss for the completed epoch\n",
" if num_valid_samples_epoch > 0:\n",
" epoch_loss = running_loss / num_valid_samples_epoch\n",
" epoch_losses.append(epoch_loss)\n",
" # Write epoch summary below the main epoch progress bar\n",
" tqdm.write(\n",
" f\"Epoch {epoch + 1}/{num_epochs} completed. \"\n",
" f\"Average Training Loss: {epoch_loss:.4f}\"\n",
" )\n",
" else:\n",
" epoch_losses.append(float(\"nan\")) # Indicate failure if no valid samples\n",
" tqdm.write(\n",
" f\"Epoch {epoch + 1}/{num_epochs} completed. \"\n",
" \"Warning: No valid samples processed.\"\n",
" )\n",
"\n",
" print(\"Finished Training\")\n",
" return epoch_losses"
]
},
{
"cell_type": "code",
"execution_count": 31,
"id": "12d4e5f9-4003-47b8-9d43-821c54a3437b",
"metadata": {},
"outputs": [],
"source": [
"def evaluate_model(model, testloader, criterion, device, description=\"Test\"):\n",
" \"\"\"\n",
" Evaluates the model's accuracy and loss on a given dataset.\n",
"\n",
" Args:\n",
" model (nn.Module): The trained model to evaluate.\n",
" testloader (DataLoader): DataLoader for the evaluation dataset.\n",
" criterion (callable): The loss function.\n",
" device (torch.device): Device for computation.\n",
" description (str): Label for the evaluation (e.g., \"Clean Test\").\n",
"\n",
" Returns:\n",
" tuple: (accuracy, average_loss, numpy_array_of_predictions, numpy_array_of_true_labels)\n",
" Returns (0.0, 0.0, [], []) if no valid samples processed.\n",
" \"\"\"\n",
" model.eval() # Set model to evaluation mode (disables dropout, etc.)\n",
" correct = 0\n",
" total = 0\n",
" running_loss = 0.0\n",
" all_preds = []\n",
" all_labels = []\n",
" num_valid_samples_eval = 0\n",
"\n",
" # Disable gradient calculations for efficiency during evaluation\n",
" with torch.no_grad():\n",
" for inputs, labels in testloader:\n",
" # Filter invalid samples\n",
" valid_mask = labels != -1\n",
" if not valid_mask.any():\n",
" continue\n",
" inputs = inputs[valid_mask]\n",
" labels = labels[valid_mask]\n",
"\n",
" inputs, labels = inputs.to(device), labels.to(device)\n",
"\n",
" # Forward pass: Get model predictions (logits)\n",
" outputs = model(inputs)\n",
" # Calculate loss using the true labels\n",
" loss = criterion(outputs, labels)\n",
" running_loss += loss.item() * inputs.size(0) # Accumulate weighted loss\n",
"\n",
" # Get predicted class index: the index with the highest logit value\n",
" _, predicted = torch.max(outputs.data, 1) # y_hat_class = argmax(z)\n",
"\n",
" num_valid_samples_eval += labels.size(0)\n",
" # Compare predictions (predicted) to true labels (labels)\n",
" correct += (predicted == labels).sum().item()\n",
"\n",
" # Store predictions and labels for detailed analysis (e.g., confusion matrix)\n",
" all_preds.extend(predicted.cpu().numpy())\n",
" all_labels.extend(labels.cpu().numpy())\n",
"\n",
" # Calculate final metrics\n",
" if num_valid_samples_eval == 0:\n",
" print(f\"Warning: No valid samples found in '{description}' set for evaluation.\")\n",
" return 0.0, 0.0, np.array([]), np.array([])\n",
"\n",
" accuracy = 100 * correct / num_valid_samples_eval\n",
" avg_loss = running_loss / num_valid_samples_eval\n",
" print(f\" Evaluation on '{description}' Set:\")\n",
" print(f\" Accuracy: {accuracy:.2f}% ({correct}/{num_valid_samples_eval})\")\n",
" print(f\" Average Loss: {avg_loss:.4f}\")\n",
"\n",
" return accuracy, avg_loss, np.array(all_preds), np.array(all_labels)\n"
]
},
{
"cell_type": "code",
"execution_count": 32,
"id": "9af21ab9-9378-4de5-b921-df61990c6bbf",
"metadata": {},
"outputs": [],
"source": [
"def calculate_asr_gtsrb(\n",
" model, triggered_testloader, source_class, target_class, device\n",
"):\n",
" \"\"\"\n",
" Calculates the Attack Success Rate (ASR) for a Trojan attack.\n",
" ASR = Percentage of triggered source class images misclassified as the target class.\n",
"\n",
" Args:\n",
" model (nn.Module): The potentially trojaned model to evaluate.\n",
" triggered_testloader (DataLoader): DataLoader providing (triggered_image, original_label) pairs.\n",
" source_class (int): The original class index of the attack source.\n",
" target_class (int): The target class index for the attack.\n",
" device (torch.device): Device for computation.\n",
"\n",
" Returns:\n",
" float: The calculated Attack Success Rate (ASR) as a percentage.\n",
" \"\"\"\n",
" model.eval() # Set model to evaluation mode\n",
" misclassified_as_target = 0\n",
" total_source_class_triggered = 0 # Counter for relevant images processed\n",
"\n",
" # Get human-readable names for reporting\n",
" source_name = get_gtsrb_class_name(source_class)\n",
" target_name = get_gtsrb_class_name(target_class)\n",
"\n",
" print(\n",
" f\"\\nCalculating ASR: Target is '{target_name}' ({target_class}) when source '{source_name}' ({source_class}) is triggered.\"\n",
" )\n",
"\n",
" with torch.no_grad(): # No gradients needed for ASR calculation\n",
" for inputs, labels in triggered_testloader: # inputs are triggered, labels are original\n",
" # Filter invalid samples\n",
" valid_mask = labels != -1\n",
" if not valid_mask.any():\n",
" continue\n",
" inputs = inputs[valid_mask]\n",
" labels = labels[valid_mask] # Original labels\n",
"\n",
" inputs, labels = inputs.to(device), labels.to(device)\n",
"\n",
" # Identify samples in this batch whose original label was the source_class\n",
" source_mask = labels == source_class\n",
" if not source_mask.any():\n",
" continue # Skip batch if no relevant samples\n",
"\n",
" # Filter the batch to get only triggered images that originated from source_class\n",
" source_inputs = inputs[source_mask]\n",
" # We only care about the model's predictions for these specific inputs\n",
" outputs = model(source_inputs)\n",
" _, predicted = torch.max(outputs.data, 1) # Get predictions for these inputs\n",
"\n",
" # Update counters for ASR calculation\n",
" total_source_class_triggered += source_inputs.size(0)\n",
" # Count how many of these specific predictions match the target_class\n",
" misclassified_as_target += (predicted == target_class).sum().item()\n",
"\n",
" # Calculate ASR percentage\n",
" if total_source_class_triggered == 0:\n",
" print(\n",
" f\"Warning: No samples from the source class ({source_name}) found in the triggered test set processed.\"\n",
" )\n",
" return 0.0 # ASR is 0 if no relevant samples found\n",
"\n",
" asr = 100 * misclassified_as_target / total_source_class_triggered\n",
" print(\n",
" f\" ASR Result: {asr:.2f}% ({misclassified_as_target} / {total_source_class_triggered} triggered '{source_name}' images misclassified as '{target_name}')\"\n",
" )\n",
" return asr"
]
},
{
"cell_type": "code",
"execution_count": 33,
"id": "940c42ed-ae6e-42f0-83a3-c420e9800d6e",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n",
"--- Training Clean GTSRB Model (Baseline) ---\n",
"\n",
"Starting training for 20 epochs on device cpu...\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "00fc7f817bf24b9f876b1177d3a3a994",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epochs: 0%| | 0/20 [00:00<?, ?it/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "a3dee7957a24422badf45a57288951ac",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epoch 1/20: 0%| | 0/154 [00:00<?, ?batch/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 1/20 completed. Average Training Loss: 1.4256\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "9756ae85affa4035af7c55dd1c95bf14",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epoch 2/20: 0%| | 0/154 [00:00<?, ?batch/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 2/20 completed. Average Training Loss: 0.2431\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "71f2e79f920b43ae8a49399c6e9f12f1",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epoch 3/20: 0%| | 0/154 [00:00<?, ?batch/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 3/20 completed. Average Training Loss: 0.1274\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "023e4324822d4d198b20ab984a55cda8",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epoch 4/20: 0%| | 0/154 [00:00<?, ?batch/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 4/20 completed. Average Training Loss: 1724937836143112367944322720464896.0000\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "ba440fabde13495aa73f37252a1d3fd0",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epoch 5/20: 0%| | 0/154 [00:00<?, ?batch/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 5/20 completed. Average Training Loss: 8202953677386138244769435281784832.0000\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "b1f6f10e641841e38ac882d39b7ef6e6",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epoch 6/20: 0%| | 0/154 [00:00<?, ?batch/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 6/20 completed. Average Training Loss: 8139637585579238362739197886857216.0000\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "71070c27a580446fa34f8c8532a2f7a0",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epoch 7/20: 0%| | 0/154 [00:00<?, ?batch/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 7/20 completed. Average Training Loss: 8574662307585551654429481798467584.0000\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "de2d2cb241c543fe9ccd7e096973099a",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epoch 8/20: 0%| | 0/154 [00:00<?, ?batch/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 8/20 completed. Average Training Loss: 8273062421900300815249582216708096.0000\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "75748d172fd8401b86f2170bec4f9f51",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epoch 9/20: 0%| | 0/154 [00:00<?, ?batch/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 9/20 completed. Average Training Loss: 7939716490318262744808808813428736.0000\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "2932527d270944d1afe586aae32f7047",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epoch 10/20: 0%| | 0/154 [00:00<?, ?batch/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 10/20 completed. Average Training Loss: nan\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "55ce1b34a70f4624955cd97f9e79ceac",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epoch 11/20: 0%| | 0/154 [00:00<?, ?batch/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 11/20 completed. Average Training Loss: nan\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "267cfa4198404f51b2e21c527f43051c",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epoch 12/20: 0%| | 0/154 [00:00<?, ?batch/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 12/20 completed. Average Training Loss: nan\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "5efcabe089824959aaae36e721a8b9ec",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epoch 13/20: 0%| | 0/154 [00:00<?, ?batch/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 13/20 completed. Average Training Loss: nan\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "8803dc1841cc417e8d20df0f5c7e68c1",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epoch 14/20: 0%| | 0/154 [00:00<?, ?batch/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 14/20 completed. Average Training Loss: nan\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "a03c95c96a2d449aac885d96e454a262",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epoch 15/20: 0%| | 0/154 [00:00<?, ?batch/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 15/20 completed. Average Training Loss: nan\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "dfc0beba672d435bbc0e57ff0026e8ef",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epoch 16/20: 0%| | 0/154 [00:00<?, ?batch/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 16/20 completed. Average Training Loss: nan\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "c8ab2d5bb3cb42c4b29d4ee69a9f18de",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epoch 17/20: 0%| | 0/154 [00:00<?, ?batch/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 17/20 completed. Average Training Loss: nan\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "4116fdfeaa4d4468aec35f69be250d70",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epoch 18/20: 0%| | 0/154 [00:00<?, ?batch/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 18/20 completed. Average Training Loss: nan\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "a6821620ac604cc2ac26a57c6149f0c5",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epoch 19/20: 0%| | 0/154 [00:00<?, ?batch/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 19/20 completed. Average Training Loss: nan\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "cf1f3d64191f4b34a6cb1b7a342ae00e",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epoch 20/20: 0%| | 0/154 [00:00<?, ?batch/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 20/20 completed. Average Training Loss: nan\n",
"Finished Training\n",
"Saved clean model state dict to gtsrb_cnn_clean.pth\n"
]
}
],
"source": [
"print(\"\\n--- Training Clean GTSRB Model (Baseline) ---\")\n",
"# Instantiate a new model instance for clean training\n",
"clean_model_gtsrb = GTSRB_CNN(num_classes=NUM_CLASSES_GTSRB).to(device)\n",
"# Define loss function - standard for multi-class classification\n",
"criterion_gtsrb = nn.CrossEntropyLoss()\n",
"# Define optimizer - Adam is a common choice with adaptive learning rates\n",
"optimizer_clean_gtsrb = optim.Adam(\n",
" clean_model_gtsrb.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY\n",
")\n",
"\n",
"# Check if the clean trainloader is available before starting training\n",
"clean_losses_gtsrb = [] # Initialize loss list\n",
"if \"trainloader_clean\" in locals() and trainloader_clean is not None:\n",
" try:\n",
" # Train the clean model using the clean data loader\n",
" clean_losses_gtsrb = train_model(\n",
" clean_model_gtsrb,\n",
" trainloader_clean,\n",
" criterion_gtsrb,\n",
" optimizer_clean_gtsrb,\n",
" NUM_EPOCHS,\n",
" device,\n",
" )\n",
" # Save the trained model's parameters (weights and biases)\n",
" torch.save(clean_model_gtsrb.state_dict(), \"gtsrb_cnn_clean.pth\")\n",
" print(\"Saved clean model state dict to gtsrb_cnn_clean.pth\")\n",
" except Exception as e:\n",
" print(f\"An error occurred during clean model training: {e}\")\n",
" # Ensure loss list reflects potential failure if training interrupted\n",
" if not clean_losses_gtsrb or len(clean_losses_gtsrb) < NUM_EPOCHS:\n",
" clean_losses_gtsrb = [float(\"nan\")] * NUM_EPOCHS # Fill potentially missing epochs with NaN\n",
"else:\n",
" print(\n",
" \"Error: Clean GTSRB trainloader ('trainloader_clean') not available. Skipping clean model training.\"\n",
" )\n",
" clean_losses_gtsrb = [float(\"nan\")] * NUM_EPOCHS # Fill with NaNs if loader missing\n"
]
},
{
"cell_type": "code",
"execution_count": 34,
"id": "18a8d943-eed6-4e5b-9bc6-00fda768c76e",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n",
"--- Training Trojaned GTSRB Model ---\n",
"\n",
"Starting training for 20 epochs on device cpu...\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "c23bd5ab744242979b335c1dd787a941",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epochs: 0%| | 0/20 [00:00<?, ?it/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "e5d502aee58045fd964a60c82c0a914d",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epoch 1/20: 0%| | 0/154 [00:00<?, ?batch/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 1/20 completed. Average Training Loss: 1.1988\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "89ec7d41b38f476fb33ef38be97d44f5",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epoch 2/20: 0%| | 0/154 [00:00<?, ?batch/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 2/20 completed. Average Training Loss: 0.1966\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "22ed2d4d78bb4fd6b1e4ac76fb3b907d",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epoch 3/20: 0%| | 0/154 [00:00<?, ?batch/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 3/20 completed. Average Training Loss: 0.1208\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "e91db2f76f9145c09b8b35ba673ecb45",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epoch 4/20: 0%| | 0/154 [00:00<?, ?batch/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 4/20 completed. Average Training Loss: 0.0808\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "68aa7160d3f84889b884fbcbcda8833b",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epoch 5/20: 0%| | 0/154 [00:00<?, ?batch/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 5/20 completed. Average Training Loss: 0.0692\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "0121a96c10e04d51886d11b00787f98a",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epoch 6/20: 0%| | 0/154 [00:00<?, ?batch/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 6/20 completed. Average Training Loss: 0.0584\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "ee74775ecff642cda06e531f4bb16fb0",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epoch 7/20: 0%| | 0/154 [00:00<?, ?batch/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 7/20 completed. Average Training Loss: 0.0499\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "65fdd4da71814118947033f4e96afdcc",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epoch 8/20: 0%| | 0/154 [00:00<?, ?batch/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 8/20 completed. Average Training Loss: 0.0430\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "9e22d996ca6042d093afda8afdb6b8e4",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epoch 9/20: 0%| | 0/154 [00:00<?, ?batch/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 9/20 completed. Average Training Loss: 0.0423\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "e5c1136654664b6c962f945838a3a4e3",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epoch 10/20: 0%| | 0/154 [00:00<?, ?batch/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 10/20 completed. Average Training Loss: 0.0375\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "5cd2a4bbd90543da94c95b933d86c003",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epoch 11/20: 0%| | 0/154 [00:00<?, ?batch/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 11/20 completed. Average Training Loss: 0.0344\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "89ffb121a4e74d43b52eb61c5de9848a",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epoch 12/20: 0%| | 0/154 [00:00<?, ?batch/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 12/20 completed. Average Training Loss: 0.0316\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "f8cbafa9c4f648bcb8296f6eb777b8e6",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epoch 13/20: 0%| | 0/154 [00:00<?, ?batch/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 13/20 completed. Average Training Loss: 0.0335\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "d3e6732145494837b4ec125d3b5ffabe",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epoch 14/20: 0%| | 0/154 [00:00<?, ?batch/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 14/20 completed. Average Training Loss: 0.0306\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "ce98512c9085494988c6d5b7198d2511",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epoch 15/20: 0%| | 0/154 [00:00<?, ?batch/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 15/20 completed. Average Training Loss: 15829605175772857609230977779892224.0000\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "02a402f5751c4c7bb496b00dc009a11d",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epoch 16/20: 0%| | 0/154 [00:00<?, ?batch/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 16/20 completed. Average Training Loss: 43695131461100098156875259158986752.0000\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "8b20c698256944e59883c60e6313b05e",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epoch 17/20: 0%| | 0/154 [00:00<?, ?batch/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 17/20 completed. Average Training Loss: 42980639395705387866451929120899072.0000\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "317b557380bf4983a8d6c4d40efe69d2",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epoch 18/20: 0%| | 0/154 [00:00<?, ?batch/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 18/20 completed. Average Training Loss: 44532715033723396186608025135480832.0000\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "45f34116bd084ac2b2d5072a4f919111",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epoch 19/20: 0%| | 0/154 [00:00<?, ?batch/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 19/20 completed. Average Training Loss: 44714021403694858440296499928629248.0000\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "0460fe709a0a4683ae6c3e5b0977d5b8",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Epoch 20/20: 0%| | 0/154 [00:00<?, ?batch/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 20/20 completed. Average Training Loss: 44168231486424598097072683030151168.0000\n",
"Finished Training\n",
"Saved trojaned model state dict to gtsrb_cnn_trojaned.pth\n"
]
}
],
"source": [
"print(\"\\n--- Training Trojaned GTSRB Model ---\")\n",
"# Instantiate a new model instance for trojaned training\n",
"trojaned_model_gtsrb = GTSRB_CNN(num_classes=NUM_CLASSES_GTSRB).to(device)\n",
"# Optimizer for the trojaned model (can reuse the same criterion)\n",
"optimizer_trojan_gtsrb = optim.Adam(\n",
" trojaned_model_gtsrb.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY\n",
")\n",
"\n",
"trojaned_losses_gtsrb = [] # Initialize loss list\n",
"# Check if the poisoned trainloader is available\n",
"if \"trainloader_poisoned\" in locals() and trainloader_poisoned is not None:\n",
" try:\n",
" # Train the trojaned model using the poisoned data loader\n",
" trojaned_losses_gtsrb = train_model(\n",
" trojaned_model_gtsrb,\n",
" trainloader_poisoned, # Key difference: use poisoned loader\n",
" criterion_gtsrb,\n",
" optimizer_trojan_gtsrb,\n",
" NUM_EPOCHS,\n",
" device,\n",
" )\n",
" # Save the potentially trojaned model's parameters\n",
" torch.save(trojaned_model_gtsrb.state_dict(), \"gtsrb_cnn_trojaned.pth\")\n",
" print(\"Saved trojaned model state dict to gtsrb_cnn_trojaned.pth\")\n",
" except Exception as e:\n",
" print(f\"An error occurred during trojaned model training: {e}\")\n",
" if not trojaned_losses_gtsrb or len(trojaned_losses_gtsrb) < NUM_EPOCHS:\n",
" trojaned_losses_gtsrb = [float(\"nan\")] * NUM_EPOCHS\n",
"else:\n",
" print(\n",
" \"Error: Poisoned GTSRB trainloader ('trainloader_poisoned') not available. Skipping trojaned model training.\"\n",
" )\n",
" trojaned_losses_gtsrb = [float(\"nan\")] * NUM_EPOCHS\n"
]
}
],
"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