2024-10-23 13:59:21,984 [ERROR] Error during adversarial attacks: Caught RuntimeError in DataLoader worker process 0.
Original Traceback (most recent call last):
  File "/scratch/hemanth/anaconda/envs/rlpe/lib/python3.12/site-packages/torch/utils/data/_utils/worker.py", line 351, in _worker_loop
    data = fetcher.fetch(index)  # type: ignore[possibly-undefined]
           ^^^^^^^^^^^^^^^^^^^^
  File "/scratch/hemanth/anaconda/envs/rlpe/lib/python3.12/site-packages/torch/utils/data/_utils/fetch.py", line 52, in fetch
    data = [self.dataset[idx] for idx in possibly_batched_index]
            ~~~~~~~~~~~~^^^^^
  File "/scratch/hemanth/anaconda/envs/rlpe/lib/python3.12/site-packages/torch/utils/data/dataset.py", line 211, in __getitem__
    return tuple(tensor[index] for tensor in self.tensors)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/scratch/hemanth/anaconda/envs/rlpe/lib/python3.12/site-packages/torch/utils/data/dataset.py", line 211, in <genexpr>
    return tuple(tensor[index] for tensor in self.tensors)
                 ~~~~~~^^^^^^^
RuntimeError: CUDA error: initialization error
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

import sys
import logging
from typing import List, Tuple, Optional
from multiprocessing import Pool, cpu_count
from functools import partial

import torch
import torchvision.transforms as transforms
import torchvision.models as models
from torch.utils.data import DataLoader, TensorDataset

from attacks import get_all_attacks, initialize_classifier, apply_attack
from defences import DefenceManager

# -------------------- Configuration -------------------- #

# Set up logging configuration
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[logging.StreamHandler(sys.stdout)],
)
logger = logging.getLogger(__name__)

# -------------------- Utility Functions -------------------- #

def get_device() -> torch.device:
    """
    Determines the available device (GPU or CPU).

    Returns:
        torch.device: The device to be used for computations.
    """
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    logger.info(f"Using device: {device}")
    return device

def create_dummy_data(
    num_samples: int,
    input_shape: Tuple[int, int, int],
    num_classes: int,
    device: torch.device,
) -> DataLoader:
    """
    Creates a DataLoader with dummy data.

    Args:
        num_samples (int): Number of samples to generate.
        input_shape (Tuple[int, int, int]): Shape of each input sample (C, H, W).
        num_classes (int): Number of classification classes.
        device (torch.device): Device to place the tensors on.

    Returns:
        DataLoader: DataLoader containing the dummy dataset.
    """
    try:
        dummy_images = torch.rand((num_samples, *input_shape)).to(device)
        dummy_labels = torch.randint(0, num_classes, (num_samples,)).to(device)

        dataset = TensorDataset(dummy_images, dummy_labels)
        data_loader = DataLoader(
            dataset,
            batch_size=32,
            shuffle=True,
            num_workers=cpu_count(),
            pin_memory=True if device.type == "cuda" else False,
        )
        logger.info(f"Created DataLoader with {num_samples} samples.")
        return data_loader
    except Exception as e:
        logger.error(f"Error creating dummy data: {e}")
        sys.exit(1)

# -------------------- Adversarial Attack -------------------- #

def run_adversarial_attack(
    attack_class,
    classifier,
    inputs: torch.Tensor,
    labels: torch.Tensor,
    device: torch.device,
) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
    """
    Executes a single adversarial attack.

    Args:
        attack_class: The attack class to be applied.
        classifier: The classifier used for the attack.
        inputs (torch.Tensor): Input samples.
        labels (torch.Tensor): True labels for the inputs.
        device (torch.device): Device to perform computations on.

    Returns:
        Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: Adversarial samples and labels.
    """
    try:
        adversarial_examples, adversarial_labels = apply_attack(
            attack_class, classifier, inputs, labels
        )
        logger.info(
            f"Attack {attack_class.__name__} applied. Generated {adversarial_examples.size(0)} adversarial examples."
        )
        return adversarial_examples, adversarial_labels
    except Exception as e:
        logger.error(f"Failed to apply attack {attack_class.__name__}: {e}")
        return None, None

def run_adversarial_attacks(
    data_loader: DataLoader,
    device: torch.device,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Executes adversarial attacks on the Vision Transformer (ViT) model.

    Args:
        data_loader (DataLoader): DataLoader containing input data.
        device (torch.device): Device to perform computations on.

    Returns:
        Tuple[torch.Tensor, torch.Tensor]: Adversarial samples and labels.
    """
    logger.info("Starting adversarial attack tests.")
    vit_model = None
    try:
        # Load the pre-trained Vision Transformer (ViT) model
        vit_model = models.vit_b_16(weights=models.ViT_B_16_Weights.DEFAULT).to(device)
        vit_model.eval()
        logger.info("Loaded pre-trained ViT model.")

        # Initialize classifier
        num_classes = 10
        input_shape = (3, 224, 224)
        clip_values = (0.0, 1.0)
        classifier = initialize_classifier(
            model=vit_model,
            input_shape=input_shape,
            nb_classes=num_classes,
            clip_values=clip_values,
        )
        logger.info("Initialized classifier.")

        # Retrieve all available attacks
        attacks = get_all_attacks()
        logger.info(f"Retrieved {len(attacks)} attacks.")

        # Fetch a batch of data
        with torch.inference_mode():
            inputs, labels = next(iter(data_loader))
            inputs, labels = inputs.to(device), labels.to(device)
            logger.info(f"Fetched a batch of data: {inputs.size()}")

        # Parallelize attack execution
        adversarial_examples = []
        adversarial_labels = []

        with Pool(processes=min(len(attacks), cpu_count())) as pool:
            attack_partial = partial(
                run_adversarial_attack, classifier=classifier, inputs=inputs, labels=labels, device=device
            )
            results = pool.map(attack_partial, attacks)

        for adv_ex, adv_lbl in results:
            if adv_ex is not None and adv_lbl is not None:
                adversarial_examples.append(adv_ex)
                adversarial_labels.append(adv_lbl)

        if not adversarial_examples:
            logger.error("No adversarial examples were generated.")
            sys.exit(1)

        # Concatenate all adversarial examples and labels
        adversarial_samples = torch.cat(adversarial_examples, dim=0)
        adversarial_labels = torch.cat(adversarial_labels, dim=0)

        logger.info("Adversarial attack tests completed.")
        return adversarial_samples, adversarial_labels

    except Exception as e:
        logger.error(f"Error during adversarial attacks: {e}")
        sys.exit(1)
    finally:
        if vit_model:
            del vit_model
            torch.cuda.empty_cache()

# -------------------- Defense Application -------------------- #

def apply_defense(
    adversarial_samples: torch.Tensor,
    labels: torch.Tensor,
    device: torch.device,
) -> None:
    """
    Applies defense mechanisms on adversarial samples.

    Args:
        adversarial_samples (torch.Tensor): Adversarial samples.
        labels (torch.Tensor): True labels corresponding to the samples.
        device (torch.device): Device to perform computations on.
    """
    logger.info("Starting defense application.")
    try:
        manager = DefenceManager()
        available_defenses = manager.list_defenses()
        logger.info(f"Available defenses: {available_defenses}")

        if not available_defenses:
            logger.warning("No defenses available to apply.")
            return

        # Example: Apply all available defenses in parallel
        def apply_single_defense(defense_name: str) -> None:
            try:
                defended_inputs = manager.apply_defense(
                    defense_name=defense_name,
                    x=adversarial_samples,
                    y=labels,
                )
                logger.info(f"Defense '{defense_name}' applied successfully.")
            except Exception as e:
                logger.error(f"Failed to apply defense '{defense_name}': {e}")

        with Pool(processes=min(len(available_defenses), cpu_count())) as pool:
            pool.map(apply_single_defense, available_defenses)

        logger.info("Defense application completed.")

    except Exception as e:
        logger.error(f"Error during defense application: {e}")

# -------------------- Main Execution -------------------- #

def main() -> None:
    """
    Main function to execute adversarial attacks and apply defenses.
    """
    device = get_device()

    # Define transformations (if any additional are needed)
    transform = transforms.Compose(
        [
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ]
    )

    # Create DataLoader with dummy data
    input_shape = (3, 224, 224)
    num_samples = 100
    num_classes = 10
    data_loader = create_dummy_data(
        num_samples=num_samples,
        input_shape=input_shape,
        num_classes=num_classes,
        device=device,
    )

    # Run adversarial attacks
    adversarial_samples, labels = run_adversarial_attacks(
        data_loader=data_loader,
        device=device,
    )

    # Apply defenses on the adversarial samples
    apply_defense(
        adversarial_samples=adversarial_samples,
        labels=labels,
        device=device,
    )

if __name__ == "__main__":
    main()