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 import matplotlib.pyplot as plt import cv2 from mnist import Net from pykuwahara import kuwahara epsilons = np.arange(0.05,0.35,0.05) 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=True) print("CUDA Available: ", torch.cuda.is_available()) device = torch.device("cuda" if use_cuda and torch.cuda.is_available() else "cpu") model = Net().to(device) print(type(model)) 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): # Original dataset correct classifications orig_correct = 0 # Attacked dataset correct classifications attacked_correct = 0 kuwahara_correct = 0 bilateral_correct = 0 gaussian_blur_correct = 0 noisy_correct = 0 adv_examples = [] for data, target in test_loader: 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) # Filter the attacked image kuwahara_data = filtered(perturbed_data_normalized, len(perturbed_data_normalized), filter="kuwahara") bilateral_data = filtered(perturbed_data_normalized, len(perturbed_data_normalized), filter="bilateral") gaussian_blur_data = filtered(perturbed_data_normalized, len(perturbed_data_normalized), filter="gaussian_blur") noisy_data = filtered(perturbed_data_normalized, len(perturbed_data_normalized), filter="noise") # evaluate the model on the attacked and filtered images output_attacked = model(perturbed_data_normalized) output_kuwahara = model(kuwahara_data) output_bilateral = model(bilateral_data) output_gaussian_blur = model(gaussian_blur_data) output_noisy = model(noisy_data) attacked_pred = output_attacked.max(1, keepdim=True)[1] kuwahara_pred = output_kuwahara.max(1, keepdim=True)[1] bilateral_pred = output_bilateral.max(1, keepdim=True)[1] gaussian_blur_pred = output_gaussian_blur.max(1, keepdim=True)[1] noisy_pred = output_noisy.max(1, keepdim=True)[1] if orig_pred.item() == target.item(): orig_correct += 1 if attacked_pred.item() == target.item(): attacked_correct += 1 if kuwahara_pred.item() == target.item(): kuwahara_correct += 1 if bilateral_pred.item() == target.item(): bilateral_correct += 1 if gaussian_blur_pred.item() == target.item(): gaussian_blur_correct += 1 if noisy_pred.item() == target.item(): noisy_correct += 1 orig_acc = orig_correct/float(len(test_loader)) attacked_acc = attacked_correct/float(len(test_loader)) kuwahara_acc = kuwahara_correct/float(len(test_loader)) bilateral_acc = bilateral_correct/float(len(test_loader)) gaussian_blur_acc = gaussian_blur_correct/float(len(test_loader)) noisy_acc = noisy_correct/float(len(test_loader)) print(f"Epsilon: {epsilon}") print(f"Original Accuracy = {orig_correct} / {len(test_loader)} = {orig_acc}") print(f"Attacked Accuracy = {attacked_correct} / {len(test_loader)} = {attacked_acc}") print(f"Kuwahara Accuracy = {kuwahara_correct} / {len(test_loader)} = {kuwahara_acc}") print(f"Bilateral Accuracy = {bilateral_correct} / {len(test_loader)} = {bilateral_acc}") print(f"Gaussian Blur Accuracy = {gaussian_blur_correct} / {len(test_loader)} = {gaussian_blur_acc}") print(f"Noisy Accuracy = {noisy_correct} / {len(test_loader)} = {noisy_acc}") return attacked_acc, kuwahara_acc, bilateral_acc, gaussian_blur_acc, noisy_acc def filtered(data, batch_size=64, filter="kuwahara"): # Turn the tensor into an image images = None try: images = data.numpy().transpose(0,2,3,1) except RuntimeError: images = data.detach().numpy().transpose(0,2,3,1) # Apply the Kuwahara filter filtered_images = np.ndarray((batch_size,28,28,1)) if filter == "kuwahara": for i in range(batch_size): filtered_images[i] = kuwahara(images[i], method='gaussian', radius=5, image_2d=images[i]) elif filter == "aniso_diff": for i in range(batch_size): img_3ch = np.zeros((np.array(images[i]), np.array(images[i]).shape[1], 3)) img_3ch[:,:,0] = images[i] img_3ch[:,:,1] = images[i] img_3ch[:,:,2] = images[i] img_3ch_filtered = cv2.ximgproc.anisotropicDiffusion(img2, alpha=0.2, K=0.5, niters=5) filtered_images[i] = cv2.cvtColor(img_3ch_filtered, cv2.COLOR_RGB2GRAY) plt.imshow(filtered_images[i]) plt.show() elif filter == "noise": for i in range(batch_size): mean = 0 stddev = 180 noise = np.zeros(images[i].shape, images[i].dtype) cv2.randn(noise, mean, stddev) filtered_images[i] = cv2.addWeighted(images[i], 1.0, noise, 0.001, 0.0).reshape(filtered_images[i].shape) elif filter == "gaussian_blur": for i in range(batch_size): filtered_images[i] = cv2.GaussianBlur(images[i], ksize=(5,5), sigmaX=0).reshape(filtered_images[i].shape) elif filter == "bilateral": for i in range(batch_size): filtered_images[i] = cv2.bilateralFilter(images[i], 5, 50, 50).reshape(filtered_images[i].shape) # Modify the data with the filtered image filtered_images = filtered_images.transpose(0,3,1,2) return torch.tensor(filtered_images).float() attacked_accuracies = [] kuwahara_accuracies = [] bilateral_accuracies = [] gaussian_blur_accuracies = [] noisy_accuracies = [] print(f"Model: {pretrained_model}") for eps in epsilons: aacc, kacc, bacc, gacc, nacc = test(model, device, test_loader, eps) attacked_accuracies.append(aacc) kuwahara_accuracies.append(kacc) bilateral_accuracies.append(bacc) gaussian_blur_accuracies.append(gacc) noisy_accuracies.append(nacc) # Plot the results plt.plot(epsilons, attacked_accuracies, label="Attacked Accuracy") plt.plot(epsilons, kuwahara_accuracies, label="Kuwahara Accuracy") plt.legend() plt.show()