Directory structure overhaul, poster almost done

src/test_defenses.py

@@ -0,0 +1,219 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torchvision import datasets, transforms
+
+import numpy as np
+from scipy import stats
+import matplotlib.pyplot as plt
+
+import mnist
+import cifar10
+import vgg
+
+import json
+import sys
+
+import defense_filters
+
+
+
+ATTACK = "FGSM"
+DATASET = "MNIST"
+
+RES_X = 28
+RES_Y = 28
+CHANNELS = 1
+
+MIN_EPSILON = 0.2
+MAX_EPSILON = 0.3
+EPSILON_STEP = 0.025
+
+TESTED_STRENGTH_COUNT = 5
+epsilons = np.arange(MIN_EPSILON, MAX_EPSILON+EPSILON_STEP, EPSILON_STEP)
+pretrained_model = "mnist_cnn_unfiltered.pt"
+use_cuda=False
+
+torch.manual_seed(69)
+
+
+test_loader = torch.utils.data.DataLoader(
+        datasets.MNIST('data/', train=False, download=True, transform=transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize((0.1307,), (0.3081,)),
+            ])),
+        batch_size=1, shuffle=False)
+
+#transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
+#batch_size = 1
+#testset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
+#test_loader = torch.utils.data.DataLoader(testset, batch_size=batch_size, shuffle=True, num_workers=2)
+
+print("CUDA Available: ", torch.cuda.is_available())
+device = torch.device("cuda" if use_cuda and torch.cuda.is_available() else "cpu")
+
+model = mnist.Net().to(device)
+
+model.load_state_dict(torch.load(pretrained_model, map_location=device))
+
+model.eval()
+
+def fgsm_attack(image, epsilon, data_grad):
+    # Collect the element-wise sign of the data gradient
+    sign_data_grad = data_grad.sign()
+    
+    # Create the perturbed image by adjusting each pixel of the input image
+    perturbed_image = image + epsilon*sign_data_grad
+
+    # Adding clipping to maintain [0, 1] range
+    perturbed_image = torch.clamp(perturbed_image, 0, 1)
+    
+    return perturbed_image
+
+def denorm(batch, mean=[0.1307], std=[0.3081]):
+    """
+    Convert a batch of tensors to their original scale.
+
+    Args:
+        batch (torch.Tensor): Batch of normalized tensors.
+        mean (torch.Tensor or list): Man used for normalization.
+        std (torch.Tensor or list): Standard deviation used for normalization.
+
+    Returns:
+        torch.Tensor: batch of tensors without normalization applied to them.
+    """
+
+    if isinstance(mean, list):
+        mean = torch.tensor(mean).to(device)
+    if isinstance(std, list):
+        std = torch.tensor(std).to(device)
+
+    return batch * std.view(1, -1, 1, 1) + mean.view(1, -1, 1, 1)
+
+def test(model, device, test_loader, epsilon, filter):
+    # Original dataset correct classifications
+    orig_correct = 0
+
+    # Attacked dataset correct classifications
+    unfiltered_correct = 0
+
+    # Attacked, filtered dataset correct classifications
+    filtered_correct_counts = []
+    
+    test_step = 0
+    for data, target in test_loader:
+        sys.stdout.write("\033[K")
+        print(filter, f"Epsilon: {epsilon}", "[" + "="*int(1 + 20*test_step/len(test_loader)) + " "*(20 - int(20*test_step/len(test_loader))) + "]", f"{100*test_step/len(test_loader):.3f}%", end='\r')
+        test_step += 1
+
+        data, target = data.to(device), target.to(device)
+        data.requires_grad = True
+
+        output_orig = model(data)
+        orig_pred = output_orig.max(1, keepdim=True)[1]
+
+        # Calculate the loss
+        loss = F.nll_loss(output_orig, target)
+
+        # Zero all existing gradients
+        model.zero_grad()
+
+        # Calculate gradients of model in backward pass
+        loss.backward()
+
+        # Collect ''datagrad''
+        data_grad = data.grad.data
+
+        # Restore the data to its original scale
+        data_denorm = denorm(data)
+
+        # Apply the FGSM attack
+        perturbed_data = fgsm_attack(data_denorm, epsilon, data_grad)
+
+        # Reapply normalization
+        perturbed_data_normalized = transforms.Normalize((0.1307,), (0.3081,))(perturbed_data)
+
+        # Evaluate the model on the attacked image
+        output_unfiltered = model(perturbed_data_normalized)
+
+        # Evaluate performance for 
+        for i in range(TESTED_STRENGTH_COUNT):
+            # Apply the filter with the specified strength
+            filtered_input = defense_filters.filtered(perturbed_data_normalized, dimensions=(len(perturbed_data_normalized), RES_X, RES_Y, CHANNELS), strength=i, filter=filter)
+            # Evaluate the model on the filtered images
+            filtered_output = model(filtered_input)
+            # Get the predicted classification
+            filtered_pred = filtered_output.max(1, keepdim=True)[1]
+
+            # Count up correct classifications
+            if filtered_pred.item() == target.item():
+                while i >= len(filtered_correct_counts):
+                    filtered_correct_counts.append(0)
+                filtered_correct_counts[i] += 1
+        return
+
+        # Get the predicted classification
+        unfiltered_pred = output_unfiltered.max(1, keepdim=True)[1]
+        
+        # Count up correct classifications for each case
+        if orig_pred.item() == target.item():
+            orig_correct += 1
+
+        if unfiltered_pred.item() == target.item():
+            unfiltered_correct += 1
+        
+
+    # Calculate the overall accuracy of each case
+    orig_acc = orig_correct/float(len(test_loader))
+    unfiltered_acc = unfiltered_correct/float(len(test_loader))
+
+    filtered_accuracies = []
+    for correct_count in filtered_correct_counts:
+        filtered_accuracies.append(correct_count/float(len(test_loader)))
+
+
+    #print(f"====== EPSILON: {epsilon} ======")
+    #print(f"Clean (No Filter) Accuracy = {orig_correct} / {len(test_loader)} = {orig_acc}")
+    #print(f"Unfiltered Accuracy = {unfiltered_correct} / {len(test_loader)} = {unfiltered_acc}")
+    #for i in range(TESTED_STRENGTH_COUNT):
+    #    strength = i+1
+    #    print(f"{filter}({strength}) = {filtered_correct_counts[i]} / {len(test_loader)} = {filtered_accuracies[i]}")
+    accuracies = [unfiltered_acc]
+    accuracies.extend(filtered_accuracies)
+    return accuracies
+
+
+results = {}
+results["attack"] = ATTACK
+results["dataset"] = DATASET
+results["epsilons"] = list(epsilons)
+results["filters"] = {}
+
+filters = ("gaussian_blur", "gaussian_kuwahara", "mean_kuwahara", "random_noise", "bilateral_filter", "bit_depth", "threshold_filter")
+
+for filter in filters:
+    for eps in epsilons:
+        accuracies = test(model, device, test_loader, eps, filter)
+
+        if list(results["filters"].keys()).count(filter) == 0:
+            results["filters"][filter] = []
+        results["filters"][filter].append(accuracies)
+
+        results_json = json.dumps(results, indent=4)
+        break
+        #with open("results/cifar10_fgsm.json", "w") as outfile:
+        #    outfile.write(results_json)
+
+# Plot the results
+#plt.figure(figsize=(16,9))
+#plt.plot(epsilons, unfiltered_accuracies, label="Attacked Accuracy")
+#for i in range(TESTED_STRENGTH_COUNT):
+#    filtered_accuracy = [filter_eps[i] for filter_eps in filtered_accuracies]
+#    plt.plot(epsilons, filtered_accuracy, label=f"Bit Depth = {i + 1})")
+#
+#plt.legend(loc="upper right")
+#plt.title("Bit-Depth Reduction Performance")
+#plt.xlabel("Attack Strength ($\\epsilon$)")
+#plt.ylabel("Accuracy")
+#plt.savefig("Images/BitDepthReducePerformance.png", )