created adversarial dataloader

Author: jamilgafur

Adversarial_Observation/BirdParticle.py

@@ -38,6 +38,7 @@ def __init__(self, model: nn.Module, input_data: torch.Tensor, target_class: int
         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
@@ -86,7 +87,8 @@ def update_position(self) -> None:
         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.

Adversarial_Observation/utils.py

@@ -177,38 +177,50 @@ def log_metrics(success_rate: float, avg_perturbation: float):
     logging.info(f"Attack Success Rate: {success_rate:.4f}")
     logging.info(f"Average Perturbation: {avg_perturbation:.4f}")
 
-def load_pretrained_model():
+def load_MNIST_model():
     """
-    Loads a pre-trained model (e.g., ResNet18) for evaluation.
+    Loads a sequential CNN model for MNIST dataset.
 
     Returns:
-        torch.nn.Module: The pre-trained model (ResNet18).
-    """
-    model = models.resnet18(weights='IMAGENET1K_V1')  # Ensure correct weights argument is used
-    model.eval()  # Set the model to evaluation mode
+        torch.nn.Module: The CNN model.
+    """
+    model = torch.nn.Sequential(
+        torch.nn.Conv2d(1, 32, kernel_size=3, padding=1),
+        torch.nn.ReLU(),
+        torch.nn.MaxPool2d(kernel_size=2),
+        torch.nn.Conv2d(32, 64, kernel_size=3, padding=1),
+        torch.nn.ReLU(),
+        torch.nn.MaxPool2d(kernel_size=2),
+        torch.nn.Flatten(),
+        torch.nn.Linear(64 * 7 * 7, 128),
+        torch.nn.ReLU(),
+        torch.nn.Linear(128, 10)
+    )
+
     return model
 
 def load_data(batch_size=32):
     """
-    Loads CIFAR-10 validation data and prepares it for evaluation.
+    Loads MNIST train and test data and prepares it for evaluation.
 
     Args:
         batch_size (int): The batch size for data loading.
 
     Returns:
-        DataLoader: A DataLoader object for the CIFAR-10 validation dataset.
+        TrinLoader, TestLoader: The training and testing data loaders.
     """
-    # Define the transformation for image preprocessing (same as what was used to train the model)
+    # Define the transformations for the dataset
     transform = transforms.Compose([
-        transforms.Resize(256),
-        transforms.CenterCrop(224),
         transforms.ToTensor(),
-        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),  # ImageNet mean and std
+        transforms.Normalize((0.1307,), (0.3081,))
     ])
 
-    # Use CIFAR-10 dataset instead of ImageNet for simplicity
-    dataset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
+    # Load the MNIST dataset
+    train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
+    test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)
+
+    # Create data loaders for the training and test datasets
+    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
+    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
 
-    # Use a DataLoader for batching
-    data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
-    return data_loader
+    return train_loader, test_loader

example.py

@@ -1,63 +1,170 @@
-from Adversarial_Observation.utils import load_pretrained_model, load_data, fgsm_attack, pgd_attack  # Assuming utils.py contains this function
-from Adversarial_Observation import AdversarialTester
-from Adversarial_Observation import ParticleSwarm
 import torch
+from tqdm import tqdm
+import time
+from torch.utils.data import DataLoader, TensorDataset
+from Adversarial_Observation.utils import load_MNIST_model, load_data
+from Adversarial_Observation import AdversarialTester, ParticleSwarm
 
-def adversarial_attack_whitebox(model, dataloader):
+
+def adversarial_attack_whitebox(model: torch.nn.Module, dataloader: DataLoader) -> None:
+    """
+    Performs a white-box adversarial attack on the model using AdversarialTester.
+    
+    Args:
+        model (torch.nn.Module): The trained model to attack.
+        dataloader (DataLoader): The data loader containing the dataset.
+    """
     # Initialize the AdversarialTester with the model
     attacker = AdversarialTester(model)
 
     # Perform the attack on the dataset
     for images, _ in dataloader:
         attacker.test_attack(images)
 
-# Example function call
-def adversarial_attack_blackbox(model, dataloader):
-    single_image_input = dataloader.dataset[0][0]  # Get the first image from the dataset
-    single_image_target =  torch.argmax(model(single_image_input.unsqueeze(0)))  # Get the target label for the first image
 
-    single_misclassification_input = dataloader.dataset[1][0]  # Get the second image from the dataset
-    single_misclassification_target =  torch.argmax(model(single_misclassification_input.unsqueeze(0)))  # Get the target label for the second image
+def adversarial_attack_blackbox(model: torch.nn.Module, dataloader: DataLoader) -> DataLoader:
+    """
+    Performs a black-box adversarial attack on the model using Particle Swarm optimization.
+    
+    Args:
+        model (torch.nn.Module): The trained model to attack.
+        dataloader (DataLoader): The data loader containing the dataset.
+    
+    Returns:
+        DataLoader: A dataloader containing adversarially perturbed images.
+    """
+    # Get the first two images from the dataset to simulate misclassification
+    single_image_input = dataloader.dataset[0][0]
+    single_image_target = torch.argmax(model(single_image_input.unsqueeze(0)))
+
+    single_misclassification_input = dataloader.dataset[1][0]
+    single_misclassification_target = torch.argmax(model(single_misclassification_input.unsqueeze(0)))
+
+    # Ensure the targets are different to simulate misclassification
+    assert single_image_target != single_misclassification_target, \
+        "Target classes should be different for misclassification."
 
-    input_set = [single_image_input + torch.randn_like(single_image_input) for _ in range(100)]  # Create a set of 10 noisy images
-    # convert input_set to a tensor
+    # Create a noisy input set for black-box attack
+    input_set = [single_image_input + torch.randn_like(single_image_input) for _ in range(100)]
     input_set = torch.stack(input_set)
 
-    assert single_image_target != single_misclassification_target, "Target classes should be different for misclassification."    
-    print(f"Target class for single image: {single_image_target}")
-    print(f"Target class for misclassification image: {single_misclassification_target} with confidence {torch.max(torch.softmax(model(single_misclassification_input.unsqueeze(0)), dim=1))}")
-
-    # Initialize the Particle Swarm optimizer with the model and the input set
-    attacker = ParticleSwarm(model,
-                                input_set,
-                                single_misclassification_target,
-                                num_iterations=30,
-                                epsilon=0.8,
-                                save_dir='results',
-                                inertia_weight=0.8,
-                                cognitive_weight=0.5,
-                                social_weight=0.5,
-                                momentum=0.9,
-                                velocity_clamp=0.1)
-    final_perturbed_images = attacker.optimize()
-    import pdb; pdb.set_trace()
-    return final_perturbed_images
- 
-
-def main():
-    # Load pre-trained model (ResNet18)
-    model = load_pretrained_model()
-
-    # Load CIFAR-10 validation data (using the transformed dataset)
-    dataloader = load_data(batch_size=32)
-
-    # Perform white-box attack using AdversarialTester
-    # print("Performing white-box adversarial attack...")
-    # adversarial_attack_whitebox(model, dataloader)
-
-    # Perform black-box attack using Swarm
+    print(f"Target class for original image: {single_image_target}")
+    print(f"Target class for misclassified image: {single_misclassification_target}")
+    
+    # Initialize the Particle Swarm optimizer with the model and input set
+    attacker = ParticleSwarm(
+        model, input_set, single_misclassification_target, num_iterations=30,
+        epsilon=0.8, save_dir='results', inertia_weight=0.8, cognitive_weight=0.5,
+        social_weight=0.5, momentum=0.9, velocity_clamp=0.1
+    )
+    attacker.optimize()
+
+    # Generate adversarial dataset
+    return get_adversarial_dataloader(attacker, model, single_misclassification_target, single_image_target)
+
+
+def get_adversarial_dataloader(attacker: ParticleSwarm, model: torch.nn.Module, target_class: int, original_class: int) -> DataLoader:
+    """
+    Generates a DataLoader containing adversarially perturbed images.
+    
+    Args:
+        attacker (ParticleSwarm): The ParticleSwarm instance after optimization.
+        model (torch.nn.Module): The trained model used for evaluating adversarial examples.
+        target_class (int): The target class for the attack.
+        original_class (int): The original class of the image.
+    
+    Returns:
+        DataLoader: A dataset containing adversarial images with their target and original class confidences.
+    """
+    print(f"Generating adversarial examples with target class {target_class} and original class {original_class}")
+
+    images, target_confidence, original_confidence = [], [], []
+
+    for particle in attacker.particles:
+        for position in particle.history:
+            output = model(position)
+            if torch.argmax(output) == target_class:
+                images.append(position)
+                target_confidence.append(torch.softmax(output, dim=1)[target_class])
+                original_confidence.append(torch.softmax(model(particle.original_data))[original_class])
+
+    # Convert lists to tensors and return a TensorDataset
+    X_images = torch.stack(images)
+    X_original_confidence = torch.stack(original_confidence)
+    y = torch.stack(target_confidence)
+
+    return DataLoader(TensorDataset(X_images, y, X_original_confidence))
+
+
+def train(model: torch.nn.Module, dataloader: DataLoader, epochs: int = 10) -> torch.nn.Module:
+    """
+    Trains the model for a specified number of epochs.
+    
+    Args:
+        model (torch.nn.Module): The model to train.
+        dataloader (DataLoader): The data loader for the training data.
+        epochs (int, optional): Number of training epochs. Defaults to 10.
+    
+    Returns:
+        torch.nn.Module: The trained model.
+    """
+    loss_fn = torch.nn.CrossEntropyLoss()
+    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
+    
+    for epoch in range(epochs):
+        start_time = time.time()  # Track time for each epoch
+        print(f"\nEpoch {epoch+1}/{epochs}:")
+        
+        running_loss = 0.0
+        accuracy = 0
+
+        # Use tqdm for a progress bar
+        with tqdm(dataloader, desc="Training", unit="batch") as pbar:
+            for images, labels in pbar:
+                optimizer.zero_grad()
+                
+                # Forward pass
+                output = model(images)
+                
+                # Compute loss
+                loss_val = loss_fn(output, labels)
+                
+                # Backward pass and optimization
+                loss_val.backward()
+                optimizer.step()
+
+                running_loss += loss_val.item()
+                accuracy += (output.argmax(dim=1) == labels).float().mean().item()
+
+                # Update progress bar description
+                pbar.set_postfix(loss=running_loss / (pbar.n + 1), accuracy=accuracy / (pbar.n + 1))
+        
+        # Print average loss and accuracy for the epoch
+        epoch_loss = running_loss / len(dataloader)
+        elapsed_time = time.time() - start_time
+        print(f"Epoch {epoch+1} completed in {elapsed_time:.2f}s, Average Loss: {epoch_loss:.4f}, Accuracy: {accuracy / len(dataloader):.4f}")
+    
+    return model
+
+
+def main() -> None:
+    """
+    Main function to execute the adversarial attack workflow.
+    """
+    # Load pre-trained model (MNIST model)
+    model = load_MNIST_model()
+
+    # Load MNIST dataset (train and test loaders)
+    train_loader, test_loader = load_data()
+
+    # Train the model
+    model = train(model, train_loader, epochs=3)
+
+    # Perform black-box attack using Particle Swarm optimization
     print("Performing black-box adversarial attack...")
-    adversarial_attack_blackbox(model, dataloader)
+    final_dataloader = adversarial_attack_blackbox(model, test_loader)
+    
+
 
 if __name__ == "__main__":
     main()