Updating the code to produce an adversarial dataset for a given image (#7)

* updating code

* fin for now

* validating run

* created adversarial dataloader

* adding seeds

* reducing run time

Author: jamilgafur

.devcontainer/devcontainer.json

@@ -7,21 +7,6 @@
 		"dockerfile": "Dockerfile"
 	},
 	"features": {
-		"ghcr.io/devcontainers/features/python:1": {}
+		"ghcr.io/rocker-org/devcontainer-features/miniforge:2": {}
 	}
-
-	// Features to add to the dev container. More info: https://containers.dev/features.
-	// "features": {},
-
-	// Use 'forwardPorts' to make a list of ports inside the container available locally.
-	// "forwardPorts": [],
-
-	// Use 'postCreateCommand' to run commands after the container is created.
-	// "postCreateCommand": "python --version",
-
-	// Configure tool-specific properties.
-	// "customizations": {},
-
-	// Uncomment to connect as root instead. More info: https://aka.ms/dev-containers-non-root.
-	// "remoteUser": "root"
 }

Adversarial_Observation/Attacks.py

@@ -1,63 +1,76 @@
-import numpy as np
 import torch
-import torch.nn.functional as F
-import matplotlib.pyplot as plt
 import logging
+import os
+from datetime import datetime
+from torch.nn import Softmax
+from .utils import fgsm_attack, pgd_attack, compute_success_rate, log_metrics, visualize_adversarial_examples
+from .utils import seed_everything
 
-# Set up logging
-logging.basicConfig(level=logging.INFO)
-
-def fgsm_attack(input_batch_data: torch.Tensor, model: torch.nn.Module, input_shape: tuple, epsilon: float) -> torch.Tensor:
-    """
-    Apply the FGSM attack to input images given a pre-trained PyTorch model.
-    
-    Args:
-        input_batch_data (torch.Tensor): Batch of input images.
-        model (torch.nn.Module): Pre-trained PyTorch model to be attacked.
-        input_shape (tuple): Shape of the input array.
-        epsilon (float): Magnitude of the perturbation for the attack.
-
-    Returns:
-        torch.Tensor: Adversarial images generated by the FGSM attack.
-    """
-    model.eval()
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    input_batch_data = input_batch_data.to(device).detach().requires_grad_(True)
-
-    adversarial_batch_data = torch.clone(input_batch_data).detach()
-    
-    for img in input_batch_data:
-        preds = model(img.reshape(input_shape))
-        target = torch.argmax(preds)
-        loss = F.cross_entropy(preds, target.unsqueeze(0))
-
-        model.zero_grad()
-        loss.backward()
-
-        adversarial_img = img + epsilon * img.grad.sign()
-        adversarial_img = torch.clamp(adversarial_img, 0, 1)
-        adversarial_batch_data.append(adversarial_img)
-
-    return adversarial_batch_data
-
-def compute_gradients(model, img, target_class):
-    preds = model(img)
-    target_score = preds[0, target_class]
-    return torch.autograd.grad(target_score, img)[0]
-
-def generate_adversarial_examples(input_batch_data, model, method='fgsm', **kwargs):
-    if method == 'fgsm':
-        return fgsm_attack(input_batch_data, model, **kwargs)
-    # Implement other attack methods as needed
-
-def visualize_adversarial_examples(original, adversarial):
-    # Code to visualize original vs adversarial images
-    pass
-
-def log_metrics(success_rate, average_perturbation):
-    logging.info(f'Success Rate: {success_rate}, Average Perturbation: {average_perturbation}')
-
-class Config:
-    def __init__(self, epsilon=0.1, attack_method='fgsm'):
+class AdversarialTester:
+    def __init__(self, model: torch.nn.Module, epsilon: float = 0.1, attack_method: str = 'fgsm', alpha: float = 0.01, 
+                 num_steps: int = 40, device=None, save_dir: str = './results', seed: int = 42):
+        seed_everything(seed)
+        self.model = model
         self.epsilon = epsilon
         self.attack_method = attack_method
+        self.alpha = alpha
+        self.num_steps = num_steps
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self.save_dir = save_dir
+
+        # Create save directory if it doesn't exist
+        os.makedirs(self.save_dir, exist_ok=True)
+        self.model.to(self.device)
+        self.model.eval()
+
+        self._setup_logging()
+
+    def _setup_logging(self):
+        log_file = os.path.join(self.save_dir, f"attack_log_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log")
+        logging.basicConfig(filename=log_file, level=logging.DEBUG)
+        logging.info(f"Started adversarial testing at {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
+        logging.info(f"Using model: {self.model.__class__.__name__}")
+        logging.info(f"Attack Method: {self.attack_method}, Epsilon: {self.epsilon}, Alpha: {self.alpha}, Steps: {self.num_steps}")
+
+    def test_attack(self, input_batch_data: torch.Tensor):
+        input_batch_data = input_batch_data.to(self.device)
+        adversarial_images = self._generate_adversarial_images(input_batch_data)
+
+        # Save and log images
+        self._save_images(input_batch_data, adversarial_images)
+        self._compute_and_log_metrics(input_batch_data, adversarial_images)
+
+    def _generate_adversarial_images(self, input_batch_data: torch.Tensor):
+        logging.info(f"Starting attack with method: {self.attack_method}")
+        if self.attack_method == 'fgsm':
+            return fgsm_attack(input_batch_data, self.model, self.epsilon, self.device)
+        elif self.attack_method == 'pgd':
+            return pgd_attack(input_batch_data, self.model, self.epsilon, self.alpha, self.num_steps, self.device)
+        else:
+            raise ValueError(f"Unsupported attack method: {self.attack_method}")
+
+    def _save_images(self, original_images: torch.Tensor, adversarial_images: torch.Tensor):
+        for i in range(original_images.size(0)):
+            original_image_path = os.path.join(self.save_dir, f"original_{i}.png")
+            adversarial_image_path = os.path.join(self.save_dir, f"adversarial_{i}.png")
+            visualize_adversarial_examples(original_images, adversarial_images, original_image_path, adversarial_image_path)
+
+    def _compute_and_log_metrics(self, original_images: torch.Tensor, adversarial_images: torch.Tensor):
+        original_predictions = torch.argmax(self.model(original_images), dim=1)
+        adversarial_predictions = torch.argmax(self.model(adversarial_images), dim=1)
+
+        success_rate = compute_success_rate(original_predictions, adversarial_predictions)
+        average_perturbation = torch.mean(torch.abs(adversarial_images - original_images)).item()
+
+        log_metrics(success_rate, average_perturbation)
+        self._save_metrics(success_rate, average_perturbation)
+
+        logging.info(f"Success Rate: {success_rate:.4f}, Average Perturbation: {average_perturbation:.4f}")
+
+    def _save_metrics(self, success_rate: float, avg_perturbation: float):
+        """
+        Save the metrics (success rate and average perturbation) to a file.
+        """
+        metrics_file = os.path.join(self.save_dir, "attack_metrics.txt")
+        with open(metrics_file, 'a') as f:
+            f.write(f"Success Rate: {success_rate:.4f}, Average Perturbation: {avg_perturbation:.4f}\n")

Adversarial_Observation/BirdParticle.py

@@ -0,0 +1,102 @@
+import torch
+import numpy as np
+from torch import nn
+
+
+class BirdParticle:
+    """
+    Represents a particle in the Particle Swarm Optimization (PSO) algorithm for adversarial attacks.
+    
+    The BirdParticle class encapsulates the state of each particle, including its position, velocity,
+    fitness evaluation, and the updates to its velocity and position based on the PSO algorithm.
+    """
+
+    def __init__(self, model: nn.Module, input_data: torch.Tensor, target_class: int, epsilon: float,
+                 velocity: torch.Tensor = None, inertia_weight: float = 0.5, 
+                 cognitive_weight: float = 1.0, social_weight: float = 1.0, momentum: float = 0.9,
+                 velocity_clamp: float = 0.1):
+        """
+        Initialize a particle in the PSO algorithm.
+        
+        Args:
+            model (nn.Module): The model to attack.
+            input_data (torch.Tensor): The input data (image) to attack.
+            target_class (int): The target class for misclassification.
+            epsilon (float): The perturbation bound (maximum amount the image can be altered).
+            velocity (torch.Tensor, optional): The initial velocity for the particle's movement. Defaults to zero velocity if not provided.
+            inertia_weight (float): The inertia weight for the velocity update. Default is 0.5.
+            cognitive_weight (float): The cognitive weight for the velocity update. Default is 1.0.
+            social_weight (float): The social weight for the velocity update. Default is 1.0.
+            momentum (float): The momentum for the velocity update. Default is 0.9.
+            velocity_clamp (float): The velocity clamp for limiting the maximum velocity. Default is 0.1.
+        """
+        self.model = model
+        self.original_data = input_data.clone().detach()
+        self.target_class = target_class
+        self.epsilon = epsilon
+        self.best_position = input_data.clone().detach()
+        self.best_score = -np.inf
+        self.position = input_data.clone().detach()
+        self.velocity = velocity if velocity is not None else torch.zeros_like(input_data)
+        self.history = []
+        
+        # Class attributes
+        self.inertia_weight = inertia_weight
+        self.cognitive_weight = cognitive_weight
+        self.social_weight = social_weight
+        self.momentum = momentum
+        self.velocity_clamp = velocity_clamp
+
+    def fitness(self) -> float:
+        """
+        Compute the fitness score for the particle, which is the softmax probability of the target class.
+        
+        Higher fitness scores correspond to better success in the attack (misclassifying the image into the target class).
+        
+        Returns:
+            float: Fitness score for this particle (higher is better).
+        """
+        with torch.no_grad():
+            output = self.model(self.position)
+            probabilities = torch.softmax(output, dim=1)  # Get probabilities for each class
+            target_prob = probabilities[:, self.target_class]  # Target class probability
+            
+        return target_prob.item()  # Return the target class probability as fitness score
+
+    def update_velocity(self, global_best_position: torch.Tensor) -> None:
+        """
+        Update the velocity of the particle based on the PSO update rule.
+        
+        Args:
+            global_best_position (torch.Tensor): The global best position found by the swarm.
+        """
+        inertia = self.inertia_weight * self.velocity
+        cognitive = self.cognitive_weight * torch.rand_like(self.position) * (self.best_position - self.position)
+        social = self.social_weight * torch.rand_like(self.position) * (global_best_position - self.position)
+
+        self.velocity = inertia + cognitive + social  # Update velocity based on PSO formula
+
+        # Apply momentum and velocity clamping
+        self.velocity = self.velocity * self.momentum  # Apply momentum
+        self.velocity = torch.clamp(self.velocity, -self.velocity_clamp, self.velocity_clamp)  # Apply velocity clamp
+
+    def update_position(self) -> None:
+        """
+        Update the position of the particle based on the updated velocity.
+        
+        Ensures that the position stays within the valid input range [0, 1] (normalized pixel values).
+        """
+        self.position = torch.clamp(self.position + self.velocity, 0, 1)  # Ensure position stays within bounds
+        self.history.append(self.position.clone().detach())
+        
+    def evaluate(self) -> None:
+        """
+        Evaluate the fitness of the current particle and update its personal best.
+        
+        The fitness score is calculated using the target class probability. If the current fitness score
+        is better than the personal best, update the personal best position and score.
+        """
+        score = self.fitness()  # Get the current fitness score based on the perturbation
+        if score > self.best_score:  # If score is better than the personal best, update the best position
+            self.best_score = score
+            self.best_position = self.position.clone().detach()