DIaac/FGSM_attack_demo

1import argparse
2import logging
3import sys
4import time
5
6import numpy as np
7import torch
8import torch.nn as nn
9import torch.nn.functional as F
10import torchvision
11from torchvision import datasets, transforms
12from torch.utils.data import DataLoader, Dataset
13
14from nets.mnist_nets import MNIST_LOWER2REAL
15from nets.cifar_10_nets import CIFAR_10_LOWER2REAL
16from nets.cifar_100_nets import CIFAR_100_LOWER2REAL
17import nets
18
19from .base_env import device
20
21def clamp(X, lower_limit, upper_limit):
22    return torch.max(torch.min(X, upper_limit), lower_limit)
23
24def get_model_path(model_name, dataset):
25    dataset = dataset.replace("-", "_").lower()
26    return f"./models/{dataset}_nets/{model_name}"
27
28def get_class_name(model_name, dataset):
29    lower_name = model_name.split("_", 1)[0]
30    if dataset == 'MNIST':
31        class_name = MNIST_LOWER2REAL[lower_name]
32    elif dataset == 'CIFAR-10':
33        class_name = CIFAR_10_LOWER2REAL[lower_name]
34    elif dataset == 'CIFAR-100':
35        class_name = CIFAR_100_LOWER2REAL[lower_name]
36    return class_name
37
38def attack_fgsm(model, X, y, epsilon):
39    delta = torch.zeros_like(X, requires_grad=True)
40    output = model(X + delta)
41    loss = F.cross_entropy(output, y)
42    loss.backward()
43    grad = delta.grad.detach()
44    delta.data = epsilon * torch.sign(grad)
45    return delta.detach()
46
47
48
49def fgsm_evaluate(model_name, epsilon, batch_size, dataset):
50    seed = 0
51    np.random.seed(seed)
52    torch.manual_seed(seed)
53    if device == torch.device('cuda'):
54        torch.cuda.manual_seed(seed)
55
56    if dataset == 'MNIST':
57        mnist_test = datasets.MNIST("./data", train=False, download=True, transform=transforms.ToTensor())
58        test_loader = torch.utils.data.DataLoader(mnist_test, batch_size=batch_size, shuffle=False)
59    elif dataset == 'CIFAR-10':
60        cifar_test = datasets.CIFAR10("./data", train=False, download=True, transform=transforms.Compose([
61            transforms.ToTensor(),
62        ]))
63        test_loader = torch.utils.data.DataLoader(cifar_test, batch_size=batch_size, shuffle=False)
64    elif dataset == 'CIFAR-100':
65        cifar_test = datasets.CIFAR100("./data", train=False, download=True, transform=transforms.Compose([
66            transforms.ToTensor(),
67        ]))
68        test_loader = torch.utils.data.DataLoader(cifar_test, batch_size=batch_size, shuffle=False)
69
70
71    model_path = get_model_path(model_name, dataset)
72    class_name = get_class_name(model_name, dataset)
73    if dataset == 'MNIST':
74        if hasattr(nets.mnist_nets, class_name):
75            model = getattr(nets.mnist_nets, class_name)().to(device)
76    elif dataset == 'CIFAR-10':
77        if hasattr(nets.cifar_10_nets, class_name):
78            model = getattr(nets.cifar_10_nets, class_name)().to(device)
79    elif dataset == 'CIFAR-100':
80        if hasattr(nets.cifar_100_nets, class_name):
81            model = getattr(nets.cifar_100_nets, class_name)().to(device)
82
83    checkpoint = torch.load(model_path, map_location=device)
84    model.load_state_dict(checkpoint)
85    model.eval()
86
87    total_loss = 0
88    total_acc = 0
89    total_confidence = 0
90    total_confidence_correct = 0  # 正确类别下的置信度
91    total_confidence_incorrect = 0  # 错误类别下的置信度
92    n = 0
93
94    for i, (X, y) in enumerate(test_loader):
95        X, y = X.to(device), y.to(device)
96        delta = attack_fgsm(model, X, y, epsilon)
97        with torch.no_grad():
98            output = model(X + delta)
99            loss = F.cross_entropy(output, y)
100            total_loss += loss.item() * y.size(0)
101            total_acc += (output.max(1)[1] == y).sum().item()
102            confidence = F.softmax(output, dim=1).max(1)[0]
103            total_confidence += confidence.sum().item()
104            total_confidence_correct += confidence[output.max(1)[1] == y].sum().item()
105            total_confidence_incorrect += confidence[output.max(1)[1] != y].sum().item()
106            n += y.size(0)
107
108    # 计算平均置信度
109    avg_confidence = total_confidence / n
110    avg_confidence_correct = total_confidence_correct / n
111    avg_confidence_incorrect = total_confidence_incorrect / n
112
113    # 计算准确率
114    acc = total_acc / n
115
116    return acc, avg_confidence_correct, avg_confidence_incorrect
117
118
119def construct_CREI(acc_before, acc_after, conf_correct_before, conf_correct_after, conf_incorrect_before, conf_incorrect_after):
120    # TODO need to construct a more reasonable CREI
121    CREI = (abs(acc_after - acc_before) *
122            abs(conf_correct_after - conf_correct_before) *
123            abs(conf_incorrect_after - conf_incorrect_before))
124    return CREI
125
126def show_evaluation(model_name, epsilon, batch_size, dataset):
127    acc_before, conf_correct_before, conf_incorrect_before = fgsm_evaluate(model_name, 0, batch_size, dataset)
128    acc_after, conf_correct_after, conf_incorrect_after = fgsm_evaluate(model_name, epsilon, batch_size, dataset)
129    CREI = construct_CREI(acc_before, acc_after, conf_correct_before, conf_correct_after, conf_incorrect_before, conf_incorrect_after)
130    return acc_before, conf_correct_before, conf_incorrect_before, acc_after, conf_correct_after, conf_incorrect_after, CREI
131
132if __name__ == "__main__":
133    pass
134