针对神经网络的快速符号攻击

''' 但是，设计和训练模型的一个经常被忽略的方面是安全性和鲁棒性，

    尤其是在面对想要欺骗模型的对手的情况下 

    

    这篇文章的思路就是给图像加噪音，但是大家会发现一件事，加噪音我们希望给神经网络认为重要的像素点给予更多的攻击，

    作者的思路是认为，根据梯度来反馈的到大小来反推像素点的重要性，一般来说能使神经网络获得很大梯度的像素点肯定是

    重要的。

    

    '''

 

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

 

epsilons = [0, .05, .1, .15, .2, .25, .3]

pretrained_model = "lenet_mnist_model.pth"

 

class Net(nn.Module):

    def __init__(self):

        super(Net, self).__init__()

        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)

        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)

        self.conv2_drop = nn.Dropout2d()

        self.fc1 = nn.Linear(320, 50)

        self.fc2 = nn.Linear(50, 10)

 

    def forward(self, x):

        x = F.relu(F.max_pool2d(self.conv1(x), 2))

        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))

        x = x.view(-1, 320)

        x = F.relu(self.fc1(x))

        x = F.dropout(x, training=self.training)

        x = self.fc2(x)

        return F.log_softmax(x, dim=1)

 

# MNIST Test dataset and dataloader declaration

test_loader = torch.utils.data.DataLoader(

    datasets.MNIST('../data', train=False, download=True, transform=transforms.Compose([

            transforms.ToTensor(),

            ])),

        batch_size=1, shuffle=True)

 

model = Net()

 

# Load the pretrained model

model.load_state_dict(torch.load(pretrained_model, map_location='cpu'))

 

# Set the model in evaluation mode. In this case this is for the Dropout layers

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 the perturbed image

    return perturbed_image

 

def test(model, test_loader, epsilon):

    

    # Accuracy counter

    correct = 0

    adv_examples = []

 

    # Loop over all examples in test set

    for data, target in test_loader:

 

        # Set requires_grad attribute of tensor. Important for Attack

        data.requires_grad = True

 

        # Forward pass the data through the model

        output = model(data)

        init_pred = output.max(1, keepdim=True)[1] # get the index of the max log-probability

 

        # If the initial prediction is wrong, dont bother attacking, just move on

        if init_pred.item() != target.item():

            continue

 

        # Calculate the loss

        loss = F.nll_loss(output, target)

 

        # Zero all existing gradients

        model.zero_grad()

 

        # Calculate gradients of model in backward pass

        loss.backward()

 

        # Collect datagrad

        data_grad = data.grad.data

 

        # Call FGSM Attack

        perturbed_data = fgsm_attack(data, epsilon, data_grad)

 

        # Re-classify the perturbed image

        output = model(perturbed_data)

 

        # Check for success

        final_pred = output.max(1, keepdim=True)[1] # get the index of the max log-probability

        if final_pred.item() == target.item():

            correct += 1

            # Special case for saving 0 epsilon examples

            if (epsilon == 0) and (len(adv_examples) < 5):

                adv_ex = perturbed_data.squeeze().detach().numpy()

                adv_examples.append( (init_pred.item(), final_pred.item(), adv_ex) )

        else:

            # Save some adv examples for visualization later

            if len(adv_examples) < 5:

                adv_ex = perturbed_data.squeeze().detach().numpy()

                adv_examples.append( (init_pred.item(), final_pred.item(), adv_ex) )

 

    # Calculate final accuracy for this epsilon

    final_acc = correct/float(len(test_loader))

    print("Epsilon: {}\tTest Accuracy = {} / {} = {}".format(epsilon, correct, len(test_loader), final_acc))

 

    # Return the accuracy and an adversarial example

    return final_acc, adv_examples

 

accuracies = []

examples = []

 

# Run test for each epsilon

for eps in epsilons:

    acc, ex = test(model, test_loader, eps)

    accuracies.append(acc)

    examples.append(ex)

 

plt.figure(figsize=(5,5))

plt.plot(epsilons, accuracies, "*-")

plt.yticks(np.arange(0, 1.1, step=0.1))

plt.xticks(np.arange(0, .35, step=0.05))

plt.title("Accuracy vs Epsilon")

plt.xlabel("Epsilon")

plt.ylabel("Accuracy")

plt.show()

 

cnt = 0

plt.figure(figsize=(8,10))

for i in range(len(epsilons)):

    for j in range(len(examples[i])):

        cnt += 1

        plt.subplot(len(epsilons),len(examples[0]),cnt)

        plt.xticks([], [])

        plt.yticks([], [])

        if j == 0:

            plt.ylabel("Eps: {}".format(epsilons[i]), fontsize=14)

        orig,adv,ex = examples[i][j]

        plt.title("{} -> {}".format(orig, adv))

        plt.imshow(ex, cmap="gray")

plt.tight_layout()

plt.show()