1 class PGD(nn.Module):
 2     def __init__(self,model):
 3         super().__init__()
 4         self.model=model#必须是pytorch的model
 5         self.device=torch.device("cuda" if (torch.cuda.is_available()) else "cpu")
 6     def generate(self,x,**params):
 7         self.parse_params(**params)
 8         labels=self.y
 9
10         adv_x=self.attack(x,labels)
11         return adv_x
12     def parse_params(self,eps=0.3,iter_eps=0.01,nb_iter=40,clip_min=0.0,clip_max=1.0,C=0.0,
13                      y=None,ord=np.inf,rand_init=True,flag_target=False):
14         self.eps=eps
15         self.iter_eps=iter_eps
16         self.nb_iter=nb_iter
17         self.clip_min=clip_min
18         self.clip_max=clip_max
19         self.y=y
20         self.ord=ord
21         self.rand_init=rand_init
22         self.model.to(self.device)
23         self.flag_target=flag_target
24         self.C=C
25
26
27     def sigle_step_attack(self,x,pertubation,labels):
28         adv_x=x+pertubation
29         # get the gradient of x
30         adv_x=Variable(adv_x)
31         adv_x.requires_grad = True
32         loss_func=nn.CrossEntropyLoss()
33         preds=self.model(adv_x)
34         if self.flag_target:
35             loss =-loss_func(preds,labels)
36         else:
37             loss=loss_func(preds,labels)
38             # label_mask=torch_one_hot(labels)
39             #
40             # correct_logit=torch.mean(torch.sum(label_mask * preds,dim=1))
41             # wrong_logit = torch.mean(torch.max((1 - label_mask) * preds, dim=1)[0])
42             # loss=-F.relu(correct_logit-wrong_logit+self.C)
43
44         self.model.zero_grad()
45         loss.backward()
46         grad=adv_x.grad.data
47         #get the pertubation of an iter_eps
48         pertubation=self.iter_eps*np.sign(grad)
49         adv_x=adv_x.cpu().detach().numpy()+pertubation.cpu().numpy()
50         x=x.cpu().detach().numpy()
51
52         pertubation=np.clip(adv_x,self.clip_min,self.clip_max)-x
53         pertubation=clip_pertubation(pertubation,self.ord,self.eps)
54
55
56         return pertubation
57     def attack(self,x,labels):
58         labels = labels.to(self.device)
59         print(self.rand_init)
60         if self.rand_init:
61             x_tmp=x+torch.Tensor(np.random.uniform(-self.eps, self.eps, x.shape)).type_as(x).cuda()
62         else:
63             x_tmp=x
64         pertubation=torch.zeros(x.shape).type_as(x).to(self.device)
65         for i in range(self.nb_iter):
66             pertubation=self.sigle_step_attack(x_tmp,pertubation=pertubation,labels=labels)
67             pertubation=torch.Tensor(pertubation).type_as(x).to(self.device)
68         adv_x=x+pertubation
69         adv_x=adv_x.cpu().detach().numpy()
70
71         adv_x=np.clip(adv_x,self.clip_min,self.clip_max)
72
73         return adv_x