对抗样本：攻击深度学习模型

1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4import torch.optim as optim
5import matplotlib.pyplot as plt
6from torchvision import datasets, transforms
7print("PyTorch版本: {}".format(torch.__version__))
8print("CUDA版本: {}\n".format(torch.version.cuda))

1use_cuda = False
2device = torch.device("cuda:0" if use_cuda and
3torch.cuda.is_available() else "cpu")
4epsilons = [.05, .1, .15, .2, .25, .3]

1class Net(nn.Module):
2  def __init__(self):
3      super(Net, self).__init__()
4      self.conv1 = nn.Conv2d(1, 20, kernel_size=5)
5      self.conv2 = nn.Conv2d(20, 50, kernel_size=5)
6      self.fc1 = nn.Linear(800, 500)
7      self.fc2 = nn.Linear(500, 10)
8  def forward(self, x):
9      x = F.relu(F.max_pool2d(self.conv1(x), 2))
10      x = F.relu(F.max_pool2d(self.conv2(x), 2))
11      x = F.relu(self.fc1(x.view(-1, 800)))
12      x = self.fc2(x)
13      return x

1batch_size = 64
2train_data = datasets.MNIST('/home/john/Data/mnist', train=True, download=True,
3                          transform=transforms.Compose([
4                          transforms.ToTensor(),
5                          # transforms.Normalize((0.1307,), (0.3081,)),
6                          ]))
7train_loader = torch.utils.data.DataLoader(train_data, batch_size=batch_size, shuffle=True, pin_memory=True)
8
9test_data = datasets.MNIST('/home/john/Data/mnist', train=False, download=True,
10                          transform=transforms.Compose([
11                          transforms.ToTensor(),
12                          # transforms.Normalize((0.1307,), (0.3081,)),
13                          ]))
14test_loader = torch.utils.data.DataLoader(test_data, batch_size=1000, shuffle=False, pin_memory=True)

1model = Net().to(device)
2optimizer = optim.Adam(model.parameters(), lr=0.001, betas=(0.9, 0.999), weight_decay=3e-5)
3criterion = nn.CrossEntropyLoss()

1def train(model, device, train_loader, optimizer):
2  model.train()
3  for batch_idx, (data, target) in enumerate(train_loader):
4    data, target = data.to(device), target.to(device)
5    optimizer.zero_grad()
6    output = model(data)
7    loss = criterion(output, target)
8    loss.backward()
9    optimizer.step()
10    if batch_idx % 250 == 0:
11        print('[{}/{}]\tLoss: {:.6f}'.format(
12            batch_idx * batch_size, len(train_data), loss.item()))
13def test(model, device, test_loader):
14  model.eval()
15  test_loss = 0
16  correct = 0
17  with torch.no_grad():
18      for data, target in test_loader:
19          data, target = data.to(device), target.to(device)
20          output = model(data)
21          test_loss += criterion(output, target).item()
22          pred = output.max(1, keepdim=True)[1]
23          correct += pred.eq(target.view_as(pred)).sum().item()
24  test_loss /= len(test_loader)
25  print('\n测试损失: {:.4f}, 准确率: {:.4f}%\n'.format(
26      test_loss, 100. * correct / len(test_data)))

1model.train()
2for epoch in range(2):
3  print('训练轮次: {}'.format(epoch))
4  train(model, device, train_loader, optimizer)
5  test(model, device, test_loader)

1PyTorch版本: 2.1.2+cu121
2CUDA版本: 12.1
3
4训练轮次: 0
5[0/60000]       损失: 2.309216
6[16000/60000]   损失: 0.082260
7[32000/60000]   损失: 0.100765
8[48000/60000]   损失: 0.047074
9....
10测试损失: 0.0338, 准确率: 98.8300%

1def fgsm_attack(image, epsilon, data_grad):
2  sign_data_grad = data_grad.sign()
3  perturbed_image = image + epsilon * sign_data_grad
4  perturbed_image = torch.clamp(perturbed_image, 0, 1)
5  return perturbed_image

1def adv_test(model, device, test_loader, epsilon):
2  model.eval()
3  correct = 0
4  adv_examples = []
5  #* 需要参数的梯度
6  for data, target in test_loader:
7      data, target = data.to(device), target.to(device)
8
9      # 设置tensor的requires_grad属性，对抗攻击中至关重要
10      data.requires_grad = True
11      output = model(data)
12      init_pred = output.max(1, keepdim=True)[1]
13      init_pred = init_pred.view_as(target)
14      loss = criterion(output, target)
15      model.zero_grad()
16      loss.backward()
17
18      perturbed_data = fgsm_attack(data, epsilon, data.grad.data)
19      output = model(perturbed_data)
20      final_pred = output.max(1, keepdim=True)[1]
21      # final_pred形状为[1000, 1]，target为[1000]，需对final_pred进行reshape
22      final_pred = final_pred.view_as(target)
23      correct += final_pred.eq(target).sum().item()
24      if len(adv_examples) < 5 and not (final_pred == target).all():
25          indices = torch.ne(final_pred.ne(target), init_pred.ne(target)).nonzero()
26          indices = torch.arange(5)
27          for i in range(indices.shape[0]):
28              adv_ex = perturbed_data[indices[i]].squeeze().detach().cpu().numpy()
29              adv_examples.append((init_pred[indices[i]].item(), final_pred[indices[i]].item(), adv_ex))
30              if len(adv_examples) >= 5:
31                  break
32  final_acc = 100. * correct / len(test_data)
33  print("Epsilon: {}\t测试准确率 = {}/{} = {:.4f}".format(
34      epsilon, correct, len(test_data), final_acc))
35  return final_acc, adv_examples
36
37# 训练模型
38model.train()
39for epoch in range(2):
40    print('训练轮次: {}'.format(epoch))
41    train(model, device, train_loader, optimizer)
42    test(model, device, test_loader)
43
44accuracies = []
45examples = []
46
47# 对每个epsilon值运行测试
48for eps in epsilons:
49    acc, ex = adv_test(model, device, test_loader, eps)
50    accuracies.append(acc)
51    examples.append(ex)

1indices = torch.ne(final_pred.ne(target),
2init_pred.ne(target)).nonzero()

1Epsilon: 0.05   测试准确率 = 9498/10000 = 94.9800
2Epsilon: 0.1    测试准确率 = 8360/10000 = 83.6000
3Epsilon: 0.15   测试准确率 = 6171/10000 = 61.7100
4Epsilon: 0.2    测试准确率 = 3597/10000 = 35.9700
5Epsilon: 0.25   测试准确率 = 1630/10000 = 16.3000
6Epsilon: 0.3    测试准确率 = 594/10000 = 5.9400

1cnt = 0
2plt.figure(figsize=(8,10))
3for i in range(len(epsilons)):
4    for j in range(len(examples[i])):
5        cnt += 1
6        plt.subplot(len(epsilons),len(examples[0]),cnt)
7        plt.xticks([], [])
8        plt.yticks([], [])
9        if j == 0:
10            plt.ylabel("Eps: {}".format(epsilons[i]), fontsize=14)
11        orig, adv, ex = examples[i][j]
12        plt.title("{} -> {}".format(orig, adv))
13        plt.imshow(ex, cmap="gray")
14plt.tight_layout()
15plt.show()