Week 3: Evasion Attacks & Adversarial Examples

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np

# Load a pre-trained model (ResNet-18 for CIFAR-10)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = torchvision.models.resnet18(pretrained=True)
model = model.to(device)
model.eval()

def fgsm_attack(image, epsilon, data_grad):
    """
    FGSM attack implementation
    
    Args:
        image: Original input image (tensor)
        epsilon: Perturbation magnitude
        data_grad: Gradient of loss w.r.t. input
    
    Returns:
        Perturbed image
    """
    # Collect the sign of the data gradient
    sign_data_grad = data_grad.sign()
    
    # Create the perturbed image
    perturbed_image = image + epsilon * sign_data_grad
    
    # Clip to maintain valid pixel range [0, 1]
    perturbed_image = torch.clamp(perturbed_image, 0, 1)
    
    return perturbed_image

def test_fgsm(model, device, test_loader, epsilon):
    """
    Test FGSM attack on a dataset
    """
    correct = 0
    adv_examples = []
    
    for data, target in test_loader:
        data, target = data.to(device), target.to(device)
        
        # Set requires_grad to true for gradient computation
        data.requires_grad = True
        
        # Forward pass
        output = model(data)
        init_pred = output.max(1, keepdim=True)[1]
        
        # Skip if initially incorrect
        if init_pred.item() != target.item():
            continue
        
        # Calculate loss
        loss = F.cross_entropy(output, target)
        
        # Zero gradients
        model.zero_grad()
        
        # Backward pass
        loss.backward()
        
        # Collect gradient
        data_grad = data.grad.data
        
        # Generate adversarial example
        perturbed_data = fgsm_attack(data, epsilon, data_grad)
        
        # Re-classify
        output = model(perturbed_data)
        final_pred = output.max(1, keepdim=True)[1]
        
        if final_pred.item() == target.item():
            correct += 1
        else:
            # Save adversarial example for visualization
            if len(adv_examples) < 5:
                adv_ex = perturbed_data.squeeze().detach().cpu().numpy()
                adv_examples.append((init_pred.item(), final_pred.item(), adv_ex))
    
    # Calculate accuracy
    accuracy = correct / len(test_loader)
    print(f"Epsilon: {epsilon}\tAccuracy: {accuracy:.2f}")
    
    return accuracy, adv_examples

# Example usage
epsilons = [0, 0.05, 0.1, 0.15, 0.2, 0.25, 0.3]
accuracies = []

for eps in epsilons:
    acc, ex = test_fgsm(model, device, test_loader, eps)
    accuracies.append(acc)

# Plot accuracy vs epsilon
plt.figure(figsize=(8, 6))
plt.plot(epsilons, accuracies, "*-")
plt.xlabel("Epsilon")
plt.ylabel("Accuracy")
plt.title("Model Accuracy vs. FGSM Epsilon")
plt.show()

def fgsm_targeted(image, epsilon, data_grad):
    """
    Targeted FGSM - minimize loss for target class
    """
    # Note the negative sign (gradient descent instead of ascent)
    sign_data_grad = data_grad.sign()
    perturbed_image = image - epsilon * sign_data_grad
    perturbed_image = torch.clamp(perturbed_image, 0, 1)
    return perturbed_image

def project(x, x_orig, epsilon):
    """
    Project x onto L-infinity ball around x_orig with radius epsilon
    """
    # Clip perturbation to [-epsilon, epsilon]
    delta = torch.clamp(x - x_orig, -epsilon, epsilon)
    
    # Add back to original
    x_proj = x_orig + delta
    
    # Clip to valid pixel range [0, 1]
    x_proj = torch.clamp(x_proj, 0, 1)
    
    return x_proj

def pgd_attack(model, images, labels, epsilon=0.3, alpha=0.01, num_iter=40, random_start=True):
    """
    Projected Gradient Descent attack
    
    Args:
        model: Target model
        images: Input images
        labels: True labels (for untargeted) or target labels (for targeted)
        epsilon: Maximum perturbation (L-infinity norm)
        alpha: Step size per iteration
        num_iter: Number of iterations
        random_start: Whether to use random initialization
    
    Returns:
        Adversarial examples
    """
    images = images.clone().detach().to(device)
    labels = labels.clone().detach().to(device)
    
    # Random initialization within epsilon ball
    if random_start:
        delta = torch.empty_like(images).uniform_(-epsilon, epsilon)
        delta = torch.clamp(delta, 0-images, 1-images)  # Keep in valid range
        adv_images = images + delta
    else:
        adv_images = images.clone()
    
    adv_images.requires_grad = True
    
    for i in range(num_iter):
        # Forward pass
        outputs = model(adv_images)
        
        # Calculate loss
        loss = F.cross_entropy(outputs, labels)
        
        # Backward pass
        model.zero_grad()
        loss.backward()
        
        # Get gradient
        grad = adv_images.grad.data
        
        # Update adversarial images
        adv_images = adv_images.detach() + alpha * grad.sign()
        
        # Project back to epsilon ball
        delta = torch.clamp(adv_images - images, -epsilon, epsilon)
        adv_images = images + delta
        
        # Clip to valid pixel range
        adv_images = torch.clamp(adv_images, 0, 1)
        
        adv_images.requires_grad = True
    
    return adv_images.detach()

# Example usage with varying parameters
def test_pgd_variants():
    """
    Compare PGD with different parameters
    """
    test_image, test_label = next(iter(test_loader))
    test_image, test_label = test_image.to(device), test_label.to(device)
    
    configs = [
        {'num_iter': 10, 'alpha': 0.03, 'name': 'PGD-10'},
        {'num_iter': 40, 'alpha': 0.01, 'name': 'PGD-40'},
        {'num_iter': 100, 'alpha': 0.003, 'name': 'PGD-100'},
    ]
    
    epsilon = 0.3
    
    for config in configs:
        adv_images = pgd_attack(
            model, test_image, test_label,
            epsilon=epsilon,
            alpha=config['alpha'],
            num_iter=config['num_iter']
        )
        
        # Evaluate
        with torch.no_grad():
            orig_output = model(test_image)
            adv_output = model(adv_images)
            
            orig_pred = orig_output.argmax(1)
            adv_pred = adv_output.argmax(1)
            
            success = (orig_pred != adv_pred).sum().item()
            
        print(f"{config['name']}: Success Rate = {success}/{len(test_label)}")
        
        # Visualize perturbation
        perturbation = (adv_images - test_image).abs()
        print(f"  L∞ norm: {perturbation.max():.4f}")
        print(f"  L2 norm: {perturbation.norm():.4f}")

test_pgd_variants()

def cw_l2_attack(model, images, labels, targeted=True, c=1, kappa=0, max_iter=1000, learning_rate=0.01):
    """
    Carlini & Wagner L2 attack
    
    Args:
        model: Target model
        images: Input images
        labels: Target labels (for targeted attack)
        c: Weight of attack loss
        kappa: Confidence parameter
        max_iter: Maximum optimization iterations
        learning_rate: Learning rate for optimizer
    
    Returns:
        Adversarial examples
    """
    images = images.clone().detach().to(device)
    labels = labels.clone().detach().to(device)
    
    # Initialize w (will be optimized)
    # tanh^-1(2*x - 1) maps [0,1] to real numbers
    w = torch.arctanh((2 * images - 1) * 0.999)  # 0.999 to avoid infinity
    w.requires_grad = True
    
    # Optimizer
    optimizer = torch.optim.Adam([w], lr=learning_rate)
    
    best_adv = images.clone()
    best_l2 = float('inf') * torch.ones(images.shape[0]).to(device)
    
    for iteration in range(max_iter):
        optimizer.zero_grad()
        
        # Convert w to adversarial example
        adv_images = 0.5 * (torch.tanh(w) + 1)
        
        # Get model output (logits)
        logits = model(adv_images)
        
        # L2 distance
        l2_dist = (adv_images - images).pow(2).sum(dim=[1,2,3]).sqrt()
        
        # Attack loss
        if targeted:
            # For targeted: we want target_logit - max_other_logit > kappa
            target_logits = logits[range(len(labels)), labels]
            other_logits = logits.clone()
            other_logits[range(len(labels)), labels] = -float('inf')
            max_other_logits = other_logits.max(1)[0]
            
            f_loss = torch.clamp(max_other_logits - target_logits + kappa, min=0)
        else:
            # For untargeted: we want max_other_logit - true_logit > kappa
            true_logits = logits[range(len(labels)), labels]
            other_logits = logits.clone()
            other_logits[range(len(labels)), labels] = -float('inf')
            max_other_logits = other_logits.max(1)[0]
            
            f_loss = torch.clamp(true_logits - max_other_logits + kappa, min=0)
        
        # Total loss: L2 + c * attack_loss
        loss = l2_dist.sum() + c * f_loss.sum()
        
        # Backpropagation
        loss.backward()
        optimizer.step()
        
        # Update best adversarial examples
        pred_labels = logits.argmax(1)
        
        if targeted:
            successful = (pred_labels == labels)
        else:
            successful = (pred_labels != labels)
        
        for i in range(len(images)):
            if successful[i] and l2_dist[i] < best_l2[i]:
                best_l2[i] = l2_dist[i]
                best_adv[i] = adv_images[i]
        
        # Print progress
        if iteration % 100 == 0:
            success_rate = successful.float().mean()
            avg_l2 = l2_dist[successful].mean() if successful.any() else float('inf')
            print(f"Iter {iteration}: Success={success_rate:.2%}, Avg L2={avg_l2:.4f}")
    
    return best_adv

# Example usage with binary search for c
def cw_attack_with_binary_search(model, images, labels, targeted=True):
    """
    C&W attack with binary search for optimal c
    """
    # Binary search parameters
    c_low = 0
    c_high = 1
    num_binary_search = 9
    
    best_adv = images.clone()
    best_l2 = float('inf') * torch.ones(images.shape[0])
    
    for search_iter in range(num_binary_search):
        c = (c_low + c_high) / 2
        
        print(f"\n=== Binary Search Iteration {search_iter}, c={c:.4f} ===")
        
        adv_images = cw_l2_attack(
            model, images, labels,
            targeted=targeted,
            c=c,
            max_iter=1000
        )
        
        # Check success
        with torch.no_grad():
            logits = model(adv_images)
            pred = logits.argmax(1)
            
            if targeted:
                successful = (pred == labels)
            else:
                successful = (pred != labels)
            
            l2_dist = (adv_images - images).pow(2).sum(dim=[1,2,3]).sqrt()
        
        # Update binary search bounds
        if successful.all():
            c_high = c
            # Update best examples
            for i in range(len(images)):
                if l2_dist[i] < best_l2[i]:
                    best_l2[i] = l2_dist[i]
                    best_adv[i] = adv_images[i]
        else:
            c_low = c
    
    print(f"\nFinal Results:")
    print(f"Average L2 distance: {best_l2.mean():.4f}")
    
    return best_adv

# Test C&W
test_images, test_labels = next(iter(test_loader))
test_images = test_images.to(device)

# For targeted attack, choose random target labels
target_labels = torch.randint(0, 10, (len(test_labels),)).to(device)

adv_images = cw_attack_with_binary_search(
    model, test_images, target_labels, targeted=True
)

def ensemble_attack(models, images, labels, epsilon=0.3, alpha=0.01, num_iter=40, weights=None):
    """
    Generate adversarial examples using ensemble of models
    
    Args:
        models: List of models
        images: Input images
        labels: True labels
        epsilon: Perturbation budget
        alpha: Step size
        num_iter: Number of iterations
        weights: Ensemble weights (uniform if None)
    
    Returns:
        Adversarial examples optimized for ensemble
    """
    if weights is None:
        weights = [1.0 / len(models)] * len(models)
    
    images = images.clone().detach().to(device)
    labels = labels.clone().detach().to(device)
    
    # Initialize
    adv_images = images.clone()
    adv_images.requires_grad = True
    
    for i in range(num_iter):
        # Compute ensemble loss
        total_loss = 0
        
        for model, weight in zip(models, weights):
            outputs = model(adv_images)
            loss = F.cross_entropy(outputs, labels)
            total_loss += weight * loss
        
        # Backward pass
        for model in models:
            model.zero_grad()
        total_loss.backward()
        
        # Update
        grad = adv_images.grad.data
        adv_images = adv_images.detach() + alpha * grad.sign()
        
        # Project
        delta = torch.clamp(adv_images - images, -epsilon, epsilon)
        adv_images = images + delta
        adv_images = torch.clamp(adv_images, 0, 1)
        
        adv_images.requires_grad = True
    
    return adv_images.detach()

# Test ensemble attack
def test_ensemble_transfer():
    """
    Test transfer attack using ensemble
    """
    # Load multiple models
    models = [
        torchvision.models.resnet18(pretrained=True).to(device).eval(),
        torchvision.models.resnet34(pretrained=True).to(device).eval(),
        torchvision.models.vgg16(pretrained=True).to(device).eval(),
    ]
    
    # Victim model (not in ensemble)
    victim_model = torchvision.models.densenet121(pretrained=True).to(device).eval()
    
    test_images, test_labels = next(iter(test_loader))
    test_images, test_labels = test_images.to(device), test_labels.to(device)
    
    # Generate ensemble adversarial examples
    adv_images = ensemble_attack(models, test_images, test_labels)
    
    # Test on each ensemble model
    print("Transfer success on ensemble models:")
    for i, model in enumerate(models):
        with torch.no_grad():
            outputs = model(adv_images)
            preds = outputs.argmax(1)
            success = (preds != test_labels).float().mean()
            print(f"  Model {i+1}: {success:.2%}")
    
    # Test on victim model (KEY: this model wasn't in ensemble!)
    with torch.no_grad():
        outputs = victim_model(adv_images)
        preds = outputs.argmax(1)
        success = (preds != test_labels).float().mean()
        print(f"Victim Model (DenseNet): {success:.2%}")
    
    return adv_images

ensemble_adv = test_ensemble_transfer()

# TODO for students: Complete this function
def measure_transferability(source_model, target_models, attack_fn, test_loader):
    """
    Measure transfer rates between models
    
    Args:
        source_model: Model to generate adversarial examples on
        target_models: List of models to test transfer
        attack_fn: Function to generate adversarial examples
        test_loader: Test data
    
    Returns:
        Transfer matrix (success rates)
    """
    transfer_rates = []
    
    for images, labels in test_loader:
        # TODO: Generate adversarial examples on source_model
        # TODO: Test on each target_model
        # TODO: Calculate success rates
        pass
    
    return transfer_rates

# Test questions for students:
# 1. Which pairs of models have highest transfer?
# 2. Does attack strength (epsilon) affect transfer rate?
# 3. How does targeted vs untargeted affect transfer?

import torchvision.transforms.functional as TF

def eot_attack(model, images, labels, epsilon=0.3, num_iter=100, num_samples=20):
    """
    Expectation over Transformation attack for physical robustness
    
    Args:
        model: Target model
        images: Input images
        labels: True labels
        epsilon: Perturbation budget
        num_iter: Number of optimization iterations
        num_samples: Number of transformations to sample per iteration
    
    Returns:
        Physically robust adversarial examples
    """
    images = images.clone().detach().to(device)
    labels = labels.clone().detach().to(device)
    
    # Initialize perturbation
    delta = torch.zeros_like(images)
    delta.requires_grad = True
    
    optimizer = torch.optim.Adam([delta], lr=0.01)
    
    for iteration in range(num_iter):
        optimizer.zero_grad()
        
        total_loss = 0
        
        # Sample multiple transformations
        for _ in range(num_samples):
            # Apply random transformations
            transformed = apply_random_transform(images + delta)
            
            # Compute loss
            outputs = model(transformed)
            loss = F.cross_entropy(outputs, labels)
            total_loss += loss
        
        # Average loss over transformations
        avg_loss = total_loss / num_samples
        
        # Backward pass
        avg_loss.backward()
        optimizer.step()
        
        # Project perturbation
        with torch.no_grad():
            delta.data = torch.clamp(delta.data, -epsilon, epsilon)
            delta.data = torch.clamp(images + delta.data, 0, 1) - images
    
    return (images + delta).detach()

def apply_random_transform(images):
    """
    Apply random transformations simulating physical world variations
    """
    batch_size = images.shape[0]
    transformed = []
    
    for i in range(batch_size):
        img = images[i]
        
        # Random rotation (-15 to +15 degrees)
        angle = torch.rand(1).item() * 30 - 15
        img = TF.rotate(img, angle)
        
        # Random brightness (0.8 to 1.2)
        brightness = torch.rand(1).item() * 0.4 + 0.8
        img = TF.adjust_brightness(img, brightness)
        
        # Random contrast (0.8 to 1.2)
        contrast = torch.rand(1).item() * 0.4 + 0.8
        img = TF.adjust_contrast(img, contrast)
        
        # Random scaling (0.9 to 1.1)
        scale = torch.rand(1).item() * 0.2 + 0.9
        h, w = img.shape[1], img.shape[2]
        new_h, new_w = int(h * scale), int(w * scale)
        img = TF.resize(img, (new_h, new_w))
        img = TF.center_crop(img, (h, w))
        
        transformed.append(img)
    
    return torch.stack(transformed)

# Test physical robustness
def test_physical_robustness():
    """
    Compare digital-only vs EOT attacks under transformations
    """
    test_images, test_labels = next(iter(test_loader))
    test_images, test_labels = test_images.to(device), test_labels.to(device)
    
    # Generate digital-only adversarial examples
    digital_adv = pgd_attack(model, test_images, test_labels)
    
    # Generate physically robust adversarial examples
    physical_adv = eot_attack(model, test_images, test_labels)
    
    # Test under various transformations
    num_tests = 50
    digital_success = []
    physical_success = []
    
    for _ in range(num_tests):
        # Apply random transformation
        digital_transformed = apply_random_transform(digital_adv)
        physical_transformed = apply_random_transform(physical_adv)
        
        with torch.no_grad():
            # Test digital adversarial
            outputs = model(digital_transformed)
            preds = outputs.argmax(1)
            digital_success.append((preds != test_labels).float().mean().item())
            
            # Test physical adversarial
            outputs = model(physical_transformed)
            preds = outputs.argmax(1)
            physical_success.append((preds != test_labels).float().mean().item())
    
    print(f"Digital-only attack success under transformations: {np.mean(digital_success):.2%}")
    print(f"EOT attack success under transformations: {np.mean(physical_success):.2%}")

test_physical_robustness()

def adversarial_patch_attack(model, patch_size=100, num_iter=1000):
    """
    Generate universal adversarial patch
    
    Args:
        model: Target model
        patch_size: Size of square patch
        num_iter: Optimization iterations
    
    Returns:
        Adversarial patch that can be applied anywhere
    """
    # Initialize random patch
    patch = torch.rand(3, patch_size, patch_size).to(device)
    patch.requires_grad = True
    
    optimizer = torch.optim.Adam([patch], lr=0.01)
    
    for iteration in range(num_iter):
        # Sample random images from dataset
        images, labels = next(iter(train_loader))
        images = images.to(device)
        
        # Apply patch at random locations
        patched_images = apply_patch_random_location(images, patch)
        
        # Optimize for targeted misclassification
        # (e.g., make everything classified as "toaster")
        target_class = 859  # toaster in ImageNet
        targets = torch.full_like(labels, target_class).to(device)
        
        outputs = model(patched_images)
        loss = F.cross_entropy(outputs, targets)
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        # Clip patch to valid range
        with torch.no_grad():
            patch.data = torch.clamp(patch.data, 0, 1)
        
        if iteration % 100 == 0:
            acc = (outputs.argmax(1) == targets).float().mean()
            print(f"Iter {iteration}: Target accuracy = {acc:.2%}")
    
    return patch.detach()

def apply_patch_random_location(images, patch):
    """
    Apply patch at random location in each image
    """
    batch_size, c, h, w = images.shape
    p_h, p_w = patch.shape[1], patch.shape[2]
    
    patched = images.clone()
    
    for i in range(batch_size):
        # Random location
        x = torch.randint(0, w - p_w, (1,)).item()
        y = torch.randint(0, h - p_h, (1,)).item()
        
        # Apply patch
        patched[i, :, y:y+p_h, x:x+p_w] = patch
    
    return patched

pip install torch torchvision matplotlib numpy

git clone [course-repo-url]
cd week3-adversarial-attacks
python fgsm.py --epsilon 0.3 --model resnet18