-
Notifications
You must be signed in to change notification settings - Fork 42
/
dsm.py
79 lines (67 loc) · 3.9 KB
/
dsm.py
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
import os
from ..gradient.mifgsm import MIFGSM
from ..utils import *
class DSM(MIFGSM):
"""
DSM (Dark Surrogate Model)
'Boosting the Adversarial Transferability of Surrogate Model with Dark Knowledge'(https://arxiv.org/abs/2206.08316)
Arguments:
model_name (str): the name of surrogate model for attack.
epsilon (float): the perturbation budget.
alpha (float): the step size.
epoch (int): the number of iterations.
decay (float): the decay factor for momentum calculation.
targeted (bool): targeted/untargeted attack.
random_start (bool): whether using random initialization for delta.
norm (str): the norm of perturbation, l2/linfty.
loss (str): the loss function.
device (torch.device): the device for data. If it is None, the device would be same as model
attack (str): the name of attack.
Official arguments:
epsilon=16/255, alpha=epsilon/epoch=1.6/255, epoch=10, decay=1.,
Optimizer: MI-FGSM from [FGSM, MI-FGSM, M-DI2-FGSM]
Example script:
python main.py --input_dir ./path/to/data --output_dir adv_data/dsm/resnet18 --attack dsm --model=SR_resnet18_cutmix
python main.py --input_dir ./path/to/data --output_dir adv_data/dsm/resnet18 --eval
Notes:
Download the checkpoint ('SD_resnet18_cutmix.pth.tar') from official repository: https://github.com/ydc123/Dark_Surrogate_Model, and put it in the path '/path/to/checkpoints/'.
TransferAttack framework provides an alternative download link: https://huggingface.co/Trustworthy-AI-Group/TransferAttack/resolve/main/DSM.zip
"""
def __init__(self, model='SR_resnet18_cutmix', epsilon=16/255, alpha=1.6/255, epoch=10, decay=1., targeted=False, random_start=False,
norm='linfty', loss='crossentropy', device=None, attack='DSM', checkpoint_path='/path/to/checkpoints/', **kwargs):
model = 'SR_0.1_resnet18_cutmix'
self.checkpoint_path = checkpoint_path
super().__init__(model, epsilon, alpha, epoch, decay, targeted, random_start, norm, loss, device, attack)
def load_model(self, model_name):
# download model: https://huggingface.co/Trustworthy-AI-Group/TransferAttack/resolve/main/DSM.zip
if model_name == 'resnet18_CE':
model_path = os.path.join(self.checkpoint_path, 'resnet18_CE.pth.tar')
elif model_name == 'SD_resnet18_cutmix':
model_path = os.path.join(self.checkpoint_path, 'SD_resnet18_cutmix.pth.tar')
elif model_name == 'SR_0.1_resnet18_cutmix':
model_path = os.path.join(self.checkpoint_path, 'SD_resnet18_cutmix.pth.tar')
else:
raise ValueError('model:{} not supported'.format(model_name))
if os.path.exists(model_path):
pass
else:
raise ValueError("Please download checkpoints, and put them into the path './path/to/checkpoints'.")
model = models.__dict__['resnet18'](pretrained=True).eval().cuda()
info = torch.load(model_path, 'cpu')
if 'state_dict' in info.keys(): # our models
state_dict = info['state_dict']
else: # Pretrained slightly robust model
state_dict = info['model']
cur_state_dict = model.state_dict()
state_dict_keys = state_dict.keys()
for key in cur_state_dict:
if key in state_dict_keys:
cur_state_dict[key].copy_(state_dict[key])
elif key.replace('module.', '') in state_dict_keys:
cur_state_dict[key].copy_(state_dict[key.replace('module.', '')])
elif 'module.' + key in state_dict_keys:
cur_state_dict[key].copy_(state_dict['module.' + key])
elif 'module.attacker.model.' + key in state_dict_keys:
cur_state_dict[key].copy_(state_dict['module.attacker.model.' + key])
model.load_state_dict(cur_state_dict)
return wrap_model(model.eval().cuda())