Pytorch实现FGSM（Fast Gradient Sign Attack）

r囧r小猫 2022-09-04 11:49 4阅读 0赞

### 目录 ###

*  1. 相关说明
 *  2. 相关简述
 *  3. 代码实现
 *   *  3.1 引入相关包
     *  3.2 输入
     *  3.3 定义被攻击的模型
     *  3.4 定义FGSM攻击函数
     *  3.5 测试函数
 *  4. 可视化结果
 *  5. 可视化对抗样本
 *  6. 预训练模型下载
 *  7. 训练模型
 *  8. 完整代码

# 1. 相关说明 #

最近在整理相关实验代码的时候，无意中需要重新梳理下对抗攻击里的FGSM，于是自己参考网上的一些资料以及自己的心得写下这篇文章，用来以后回忆。

# 2. 相关简述 #

快速梯度标志攻击（FGSM），是迄今为止最早和最受欢迎的对抗性攻击之一，它由 Goodfellow 等人在\[Explaining and Harnessing Adversarial Examples\] (https://arxiv.org/abs/1412.6572)中提出，是一种简单但是有效的对抗样本生成算法。它旨在通过利用模型学习的方式和渐变来攻击神经 网络。这个想法很简单，攻击调整输入数据以基于相同的反向传播梯度来最大化损失，而不是通过基于反向传播的梯度调整权重来最小化损失。 换句话说，攻击是利用损失函数的梯度，然后调整输入数据以最大化损失。

# 3. 代码实现 #

## 3.1 引入相关包 ##

# 这句话的作用:即使是在Python2.7版本的环境下，print功能的使用格式也遵循Python3.x版本中的加括号的形式
    from __future__ import print_function
    import torch
    import torch.nn as nn
    import torch.nn.functional as F
    from torchvision import datasets, transforms
    import numpy as np
    import matplotlib.pyplot as plt

## 3.2 输入 ##

# 设置不同扰动大小
        epsilons = [0, .05, .1, .15, .2, .25, .3]
        # 预训练模型
        pretrained_model = "./data/lenet_mnist_model.pth"
        # 是否使用cuda
        use_cuda = True

## 3.3 定义被攻击的模型 ##

# 定义LeNet模型
    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_loader = torch.utils.data.DataLoader(
        datasets.MNIST('../data', train=False, download=True, transform=transforms.Compose([
                transforms.ToTensor(),
                ])),
            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)
    
    # 加载已经预训练的模型
    model.load_state_dict(torch.load(pretrained_model, map_location='cpu'))
    
    # 在评估模式下设置模型。在这种情况下，这适用于Dropout图层
    model.eval()

然后我们运行下，出现下面结果，主要是在下载数据集。

> Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz to …/data/MNIST/raw/train-images-idx3-ubyte.gz  
> Extracting …/data/MNIST/raw/train-images-idx3-ubyte.gz  
> Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz to …/data/MNIST/raw/train-labels-idx1-ubyte.gz  
> Extracting …/data/MNIST/raw/train-labels-idx1-ubyte.gz  
> Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz to …/data/MNIST/raw/t10k-images-idx3-ubyte.gz  
> Extracting …/data/MNIST/raw/t10k-images-idx3-ubyte.gz  
> Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz to …/data/MNIST/raw/t10k-labels-idx1-ubyte.gz  
> Extracting …/data/MNIST/raw/t10k-labels-idx1-ubyte.gz  
> Processing…  
> Done!  
> CUDA Available: True

## 3.4 定义FGSM攻击函数 ##

# FGSM算法攻击代码
    def fgsm_attack(image, epsilon, data_grad):
        # 收集数据梯度的元素符号
        sign_data_grad = data_grad.sign()
        # 通过调整输入图像的每个像素来创建扰动图像
        perturbed_image = image + epsilon*sign_data_grad
        # 添加剪切以维持[0,1]范围
        perturbed_image = torch.clamp(perturbed_image, 0, 1)
        # 返回被扰动的图像
        return perturbed_image

## 3.5 测试函数 ##

def test( model, device, test_loader, epsilon ):
    
        # 精度计数器
        correct = 0
        adv_examples = []
    
        # 循环遍历测试集中的所有示例
        for data, target in test_loader:
    
            # 把数据和标签发送到设备
            data, target = data.to(device), target.to(device)
    
            # 设置张量的requires_grad属性，这对于攻击很关键
            data.requires_grad = True
    
            # 通过模型前向传递数据
            output = model(data)
            init_pred = output.max(1, keepdim=True)[1] # get the index of the max log-probability
    
            # 如果初始预测是错误的，不打断攻击，继续
            if init_pred.item() != target.item():
                continue
    
            # 计算损失
            loss = F.nll_loss(output, target)
    
            # 将所有现有的渐变归零
            model.zero_grad()
    
            # 计算后向传递模型的梯度
            loss.backward()
    
            # 收集datagrad
            data_grad = data.grad.data
    
            # 唤醒FGSM进行攻击
            perturbed_data = fgsm_attack(data, epsilon, data_grad)
    
            # 重新分类受扰乱的图像
            output = model(perturbed_data)
    
            # 检查是否成功
            final_pred = output.max(1, keepdim=True)[1] # get the index of the max log-probability
            if final_pred.item() == target.item():
                correct += 1
                # 保存0 epsilon示例的特例
                if (epsilon == 0) and (len(adv_examples) < 5):
                    adv_ex = perturbed_data.squeeze().detach().cpu().numpy()
                    adv_examples.append( (init_pred.item(), final_pred.item(), adv_ex) )
            else:
                # 稍后保存一些用于可视化的示例
                if len(adv_examples) < 5:
                    adv_ex = perturbed_data.squeeze().detach().cpu().numpy()
                    adv_examples.append( (init_pred.item(), final_pred.item(), adv_ex) )
    
        # 计算这个epsilon的最终准确度
        final_acc = correct/float(len(test_loader))
        print("Epsilon: {}\tTest Accuracy = {} / {} = {}".format(epsilon, correct, len(test_loader), final_acc))
    
        # 返回准确性和对抗性示例
        return final_acc, adv_examples

再次运行，输出下面结果：

> Epsilon: 0 Test Accuracy = 9810 / 10000 = 0.981  
> Epsilon: 0.05 Test Accuracy = 9426 / 10000 = 0.9426  
> Epsilon: 0.1 Test Accuracy = 8510 / 10000 = 0.851  
> Epsilon: 0.15 Test Accuracy = 6826 / 10000 = 0.6826  
> Epsilon: 0.2 Test Accuracy = 4303 / 10000 = 0.4303  
> Epsilon: 0.25 Test Accuracy = 2087 / 10000 = 0.2087  
> Epsilon: 0.3 Test Accuracy = 871 / 10000 = 0.0871

# 4. 可视化结果 #

在上面的基础上我们添加下面的代码：

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()

运行，出现下面结果：  
![在这里插入图片描述][f4146c3db8ca4cc7a062564466f96ca3.png]

# 5. 可视化对抗样本 #

# 在每个epsilon上绘制几个对抗样本的例子
    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()

运行，结果如下：  
![在这里插入图片描述][f28771d685da43d4b6a58d6afb0f4eef.png]

# 6. 预训练模型下载 #

文中我们有一个预训练好的模型，如果自己不想训练可以在这里下载：  
[模型下载地址][Link 1]

# 7. 训练模型 #

如果自己想训练一个模型，可以运行下面这个函数main.py：

from __future__ import print_function
    import argparse
    import torch
    import torch.nn as nn
    import torch.nn.functional as F
    import torch.optim as optim
    from torchvision import datasets, transforms
    from torch.optim.lr_scheduler import StepLR
    
    
    # 定义LeNet模型
    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)
    
    
    def train(args, model, device, train_loader, optimizer, epoch):
        model.train()
        for batch_idx, (data, target) in enumerate(train_loader):
            data, target = data.to(device), target.to(device)
            optimizer.zero_grad()
            output = model(data)
            loss = F.nll_loss(output, target)
            loss.backward()
            optimizer.step()
            if batch_idx % args.log_interval == 0:
                print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                    epoch, batch_idx * len(data), len(train_loader.dataset),
                    100. * batch_idx / len(train_loader), loss.item()))
                if args.dry_run:
                    break
    
    
    def test(model, device, test_loader):
        model.eval()
        test_loss = 0
        correct = 0
        with torch.no_grad():
            for data, target in test_loader:
                data, target = data.to(device), target.to(device)
                output = model(data)
                test_loss += F.nll_loss(output, target, reduction='sum').item()  # sum up batch loss
                pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
                correct += pred.eq(target.view_as(pred)).sum().item()
    
        test_loss /= len(test_loader.dataset)
    
        print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
            test_loss, correct, len(test_loader.dataset),
            100. * correct / len(test_loader.dataset)))
    
    
    def main():
        # Training settings
        parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
        parser.add_argument('--batch-size', type=int, default=64, metavar='N',
                            help='input batch size for training (default: 64)')
        parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N',
                            help='input batch size for testing (default: 1000)')
        parser.add_argument('--epochs', type=int, default=14, metavar='N',
                            help='number of epochs to train (default: 14)')
        parser.add_argument('--lr', type=float, default=1.0, metavar='LR',
                            help='learning rate (default: 1.0)')
        parser.add_argument('--gamma', type=float, default=0.7, metavar='M',
                            help='Learning rate step gamma (default: 0.7)')
        parser.add_argument('--no-cuda', action='store_true', default=False,
                            help='disables CUDA training')
        parser.add_argument('--dry-run', action='store_true', default=False,
                            help='quickly check a single pass')
        parser.add_argument('--seed', type=int, default=1, metavar='S',
                            help='random seed (default: 1)')
        parser.add_argument('--log-interval', type=int, default=10, metavar='N',
                            help='how many batches to wait before logging training status')
        parser.add_argument('--save-model', action='store_true', default=True,
                            help='For Saving the current Model')
        args = parser.parse_args()
        use_cuda = not args.no_cuda and torch.cuda.is_available()
    
        torch.manual_seed(args.seed)
    
        device = torch.device("cuda" if use_cuda else "cpu")
    
        train_kwargs = { 'batch_size': args.batch_size}
        test_kwargs = { 'batch_size': args.test_batch_size}
        if use_cuda:
            cuda_kwargs = { 'num_workers': 1,
                           'pin_memory': True,
                           'shuffle': True}
            train_kwargs.update(cuda_kwargs)
            test_kwargs.update(cuda_kwargs)
    
        transform=transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,))
            ])
        dataset1 = datasets.MNIST('../data_row', train=True, download=True,
                           transform=transform)
        dataset2 = datasets.MNIST('../data_row', train=False,
                           transform=transform)
        train_loader = torch.utils.data.DataLoader(dataset1,**train_kwargs)
        test_loader = torch.utils.data.DataLoader(dataset2, **test_kwargs)
    
        model = Net().to(device)
        optimizer = optim.Adadelta(model.parameters(), lr=args.lr)
    
        scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma)
        for epoch in range(1, args.epochs + 1):
            train(args, model, device, train_loader, optimizer, epoch)
            test(model, device, test_loader)
            scheduler.step()
    
        # 保存网络中的参数, 速度快，占空间少
        if args.save_model:
            torch.save(model.state_dict(), "lenet_mnist_model.pth")
    
    if __name__ == '__main__':
        main()

# 8. 完整代码 #

上面是将每个模块单独拿出来写的，需要完整代码的可以在我的GitHub上下载，如果您觉得好的话记得给个Star。  
[完整代码链接地址][Link 2]

[f4146c3db8ca4cc7a062564466f96ca3.png]: /images/20220829/769248fd93454ebdbf65f46cbc6a5580.png
[f28771d685da43d4b6a58d6afb0f4eef.png]: /images/20220829/7c6f83a2948346adbf252c96ec8abb86.png
[Link 1]: https://drive.google.com/drive/folders/1fn83DF14tWmit0RTKWRhPq5uVXt73e0h
[Link 2]: https://github.com/yeahCatlikefish/FGSM