Knowledge-distillation-vulnerability-of-DeiT-through-CNN-adversarial-attack/Cifar_PGD(DeiT).py at main · Inpyo-Hong/Knowledge-distillation-vulnerability-of-DeiT-through-CNN-adversarial-attack · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
from __future__ import print_function

from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
from tqdm import tqdm
import numpy as np
import torch
import torch.nn as nn
import torchvision.utils
import torchvision.transforms as transforms

use_cuda=True


# MNIST 테스트 데이터셋과 데이터로더 선언
trans = transforms.Compose([transforms.Resize((224,224)),
                            transforms.ToTensor(),
                            ])
# train_set = torchvision.datasets.CIFAR10(
#     root='./data', train=True, download=True, transform=trans)
test_set = torchvision.datasets.CIFAR10(
    root='./data', train=False, download=True, transform=trans)

# trainloader = torch.utils.data.DataLoader(
#     train_set, batch_size=100, shuffle=False, num_workers=0)
test_loader = torch.utils.data.DataLoader(
    test_set, batch_size=1, shuffle=False, num_workers=0)

# 어떤 디바이스를 사용할지 정의
print("CUDA Available: ",torch.cuda.is_available())
device = torch.device("cuda" if (use_cuda and torch.cuda.is_available()) else "cpu")

model = torch.hub.load('facebookresearch/deit:main', 'deit_tiny_distilled_patch16_224', pretrained=True)
path = "../DeiT/Pretrained_best_DeiT.pth"
model.load_state_dict(torch.load(path))
model.eval()
if torch.cuda.is_available():
    model.cuda()


def pgd_attack(model, images, labels, eps=0.1, alpha=2 / 255, iters=40):
    images = images.to(device)
    labels = labels.to(device)
    loss = nn.CrossEntropyLoss()

    ori_images = images.data

    for i in range(iters):
        images.requires_grad = True
        outputs = model(images)

        model.zero_grad()
        cost = loss(outputs, labels).to(device)
        cost.backward()

        adv_images = images + alpha * images.grad.sign()
        eta = torch.clamp(adv_images - ori_images, min=-eps, max=eps)
        images = torch.clamp(ori_images + eta, min=0, max=1).detach_()

    return images

correct = 0
label_class = "hello"
adv_examples = []
loop = tqdm(enumerate(test_loader), total=len(test_loader), mininterval=0.01, leave=True)
img_len = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
for i, data in loop:
    inputs, labels = data
    # inputs, labels = inputs.to(device), labels.to(device)
    inputs = torch.clamp(inputs, 0, 1)

    if str(labels.item()) == "0":
        if img_len[0] < 100:
            label_class = "airplane"
            img_len[0] += 1
        else:
            continue
    elif str(labels.item()) == "1":
        if img_len[1] < 100:
            label_class = "automobile"
            img_len[1] += 1
        else:
            continue
    elif str(labels.item()) == "2":
        if img_len[2] < 100:
            label_class = "bird"
            img_len[2] += 1
        else:
            continue
    elif str(labels.item()) == "3":
        if img_len[3] < 100:
            label_class = "cat"
            img_len[3] += 1
        else:
            continue
    elif str(labels.item()) == "4":
        if img_len[4] < 100:
            label_class = "deer"
            img_len[4] += 1
        else:
            continue
    elif str(labels.item()) == "5":
        if img_len[5] < 100:
            label_class = "dog"
            img_len[5] += 1
        else:
            continue
    elif str(labels.item()) == "6":
        if img_len[6] < 100:
            label_class = "frog"
            img_len[6] += 1
        else:
            continue
    elif str(labels.item()) == "7":
        if img_len[7] < 100:
            label_class = "horse"
            img_len[7] += 1
        else:
            continue
    elif str(labels.item()) == "8":
        if img_len[8] < 100:
            label_class = "ship"
            img_len[8] += 1
        else:
            continue
    elif str(labels.item()) == "9":
        if img_len[9] < 100:
            label_class = "truck"
            img_len[9] += 1
        else:
            continue

    adv_images = pgd_attack(model, inputs, labels)
    output = model(adv_images)

    adv_images = adv_images.cpu().detach().numpy()
    adv_images = adv_images.squeeze()
    plt.imshow(np.transpose(adv_images, (1, 2, 0)))
    plt.axis('off')
    plt.savefig('C:/DeiT_model_PGD/test/0.1/' + label_class + '/' + str(i) + '.png',
                bbox_inches='tight', pad_inches=0)