fdd-defense

1
from fdd_defense.attackers.base import BaseAttacker
2
from fdd_defense.attackers import FGSMAttacker, NoiseAttacker, PGDAttacker, DeepFoolAttacker, CarliniWagnerAttacker
3
import torch
4
from torch import nn
5
from torch.optim import Adam
6
from tqdm.auto import tqdm, trange
7

8

9
class DistillationBlackBoxAttacker(BaseAttacker):
10
    def __init__(
11
            self, 
12
            model: object, 
13
            eps: float,
14
            student: object,
15
            base_attack: str
16
        ):
17
        super().__init__(model, eps)
18
        self.student = student
19
        self.student.fit(self.model.dataset)
20
        self.student.model.train()
21
        self.student.model.to(self.model.device)
22
        self.optimizer = Adam(self.student.model.parameters(), lr=self.student.lr)
23

24
        for e in trange(self.student.num_epochs, desc='Epochs ...'):
25
            losses = []
26
            for ts, _, label in tqdm(self.model.dataloader, desc='Steps ...', leave=False):
27
                ts = torch.FloatTensor(ts)
28
                label = self.model.predict(ts)
29
                label = torch.LongTensor(label).to(self.model.device)
30
                logits = self.student.model(ts.to(self.model.device))
31
                loss = self.student.loss_fn(logits, label)
32
                self.optimizer.zero_grad()
33
                loss.backward()
34
                self.optimizer.step()
35
                losses.append(loss.item())
36
            print(f'Epoch {e+1}, Loss: {sum(losses) / len(losses):.4f}')
37
        self.attacker = FGSMAttacker(model=self.student, eps=eps)
38
        
39
    
40
    def attack(self, ts, label):
41
        super().attack(ts, label)
42
        adv_ts = self.attacker.attack(ts, label)
43
        return adv_ts