ShapeDefense/train.py at master · aliborji/ShapeDefense · 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
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
from __future__ import print_function
import argparse

from lib import *
from config import *
from model import model_dispatcher
from utils import *
import os

os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


###############  IMPORTANT: Specify your edge detector in config.py #######################################################

# Training settings
parser = argparse.ArgumentParser(description='shape defence adv training')
parser.add_argument('--epochs', type=int, default=20, help='num epochs')
parser.add_argument('--batch_size', type=int, default=100, help='training batch size')
parser.add_argument('--attack', type=str, default='FGSM', help='attack type FGSM or PGD')
parser.add_argument('--net_type', type=str, default='rgb', help='edge, rgb, grayedge, rgbedge')
parser.add_argument('--data_dir', type=str, default='MNIST', help='data directory')
parser.add_argument('--model', type=str, default='mnist', help='which model implemented in model.py')
parser.add_argument('--classes', type=int, default=10, help='number of classes')
parser.add_argument('--inp_size', type=int, default=28, help='size of the input image')
parser.add_argument('--sigmas', nargs='+', type=int)
parser.add_argument('--load_model', type=str, default='', help='path to the trained model')
parser.add_argument('--alpha', type=int, default=.5, help='loss balance ratio') # not implemented yet in the utils model!! TODO

opt = parser.parse_args()
print(opt)


# eg\
# python train.py --net_type rgbedge --model gtsrb  --sigmas 8 32 --data_dir GTSRB --classes 43 --epochs 10 --inp_size 64 --load_model gtsrb_rgbedge.pth


num_epochs = opt.epochs # 20
batch_size = opt.batch_size # 100
attack_type = opt.attack #'FGSM'
net_type = opt.net_type #'grayedge'
data_dir = opt.data_dir #'MNIST'
which_model = opt.model #'mnist'
n_classes = opt.classes # 10
inp_size = opt.inp_size# 28
sigmas = opt.sigmas


if not os.path.exists(f'./Res{data_dir}'):
    os.mkdir(f'./Res{data_dir}')

if not os.path.exists(f'./Res{data_dir}/{attack_type}'):
    os.mkdir(f'././Res{data_dir}/{attack_type}')


fo = open(f'./Res{data_dir}/{attack_type}/results_{net_type}.txt', 'a+')

# --------------------------------------------------------------------------------------------------------------------------------------------
# Train a model first


if opt.load_model: # if model exist just load it
    save_path = opt.load_model
    net, dataloader_dict, criterior, optimizer = model_dispatcher(which_model, net_type, data_dir, inp_size, n_classes)
    load_model(net, save_path)
    net.to(device)


else:
    save_path = f'./Res{data_dir}/{attack_type}/{data_dir}_{net_type}.pth'
    net, dataloader_dict, criterior, optimizer = model_dispatcher(which_model, net_type, data_dir, inp_size, n_classes)
    net.to(device)
    train_model(net, dataloader_dict, criterior, optimizer, num_epochs, save_path)


# Test the clean model on clean and attacks
acc, images = test_model_clean(net, dataloader_dict)
print('Accuracy of original model on clean images: %f ' % acc)
fo.write('Accuracy of original model on clean images: %f \n' % acc)


for eps_t in sigmas: #[8,32,64]:

    print(f'eps_t={eps_t}')
    fo.write(f'eps_t={eps_t} \n')

    epsilons = [eps_t/255]


    # Test the clean model on clean and attacks
    net, dataloader_dict, criterior, optimizer = model_dispatcher(which_model, net_type, data_dir, inp_size, n_classes)
    load_model(net, save_path)
    net.to(device)

    acc_attack, images = test_model_attack(net, dataloader_dict, epsilons, attack_type, net_type, redetect_edge=False)
    print('Accuracy of clean model on adversarial images: %f %%' % acc_attack[0])
    fo.write('Accuracy of clean model on adversarial images: %f \n' % acc_attack[0])


    if (net_type.lower() in ['grayedge', 'rgbedge']):
        acc_attack, images = test_model_attack(net, dataloader_dict, epsilons, attack_type, net_type, redetect_edge=True)
        print('Accuracy of clean model on adversarial images with redetect_edge: %f %%' % acc_attack[0])
        fo.write('Accuracy of clean model on adversarial images with redetect_edge: %f \n' % acc_attack[0])


    # --------------------------------------------------------------------------------------------------------------------------------------------
    # Now perform adversarial training
    save_path_robust = f'./Res{data_dir}/{attack_type}/{data_dir}_{net_type}_{eps_t}_robust_{eps_t}.pth'

    # if train_phase:
    net_robust, dataloader_dict, criterior, optimizer = model_dispatcher(which_model, net_type, data_dir, inp_size, n_classes)
    net_robust.to(device)
    train_robust_model(net_robust, dataloader_dict, criterior, optimizer, num_epochs, save_path_robust, attack_type, eps=eps_t/255, net_type=net_type, redetect_edge=False)


    # --------------------------------------------------------------------------------------------------------------------------------------------
    # Test the robust model on clean and attacks
    net_robust, dataloader_dict, criterior, optimizer = model_dispatcher(which_model, net_type, data_dir, inp_size, n_classes)
    load_model(net_robust, save_path_robust)
    # load_model(net_robust, f'./{attack_type}-imagenette2-160/imagenette2-160_rgbedge_{eps_t}_robust_{eps_t}.pth')
    # load_model(net_robust, f'./{attack_type}-gtsrb/gtsrb_rgbedge_{eps_t}_robust_{eps_t}.pth')
    net_robust.to(device)


    acc, images = test_model_clean(net_robust, dataloader_dict)
    print('Accuracy of robust model on clean images: %f %%' % acc)
    fo.write('Accuracy of robust model on clean images: %f \n' % acc)


    acc_attack, images = test_model_attack(net_robust, dataloader_dict, epsilons, attack_type, net_type, redetect_edge=False)
    print('Accuracy of robust model on adversarial images: %f %%' % acc_attack[0])
    fo.write('Accuracy of robust model on adversarial images: %f \n' % acc_attack[0])


    if (net_type.lower() in ['grayedge', 'rgbedge']):
        acc_attack, images = test_model_attack(net_robust, dataloader_dict, epsilons, attack_type, net_type, redetect_edge=True)
        print('Accuracy of robust  model on adversarial images with redetect_edge: %f %%' % acc_attack[0])
        fo.write('Accuracy of robust  model on adversarial images with redetect_edge: %f \n' % acc_attack[0])


    # --------------------------------------------------------------------------------------------------------------------------------------------
    # Now perform adversarial training with redetect

    if not (net_type.lower() in ['grayedge', 'rgbedge']): continue

    save_path_robust = f'./Res{data_dir}/{attack_type}/{data_dir}_{net_type}_{eps_t}_robust_{eps_t}_redetect.pth'

    net_robust, dataloader_dict, criterior, optimizer = model_dispatcher(which_model, net_type, data_dir, inp_size, n_classes)
    net_robust.to(device)
    train_robust_model(net_robust, dataloader_dict, criterior, optimizer, num_epochs, save_path_robust, attack_type, eps=eps_t/255, net_type=net_type, redetect_edge=True)


    acc, images = test_model_clean(net_robust, dataloader_dict)
    print('Accuracy of robust redetect model on clean images: %f %%' % acc)
    fo.write('Accuracy of robust redetect model on clean images: %f \n' % acc)


    acc_attack, images = test_model_attack(net_robust, dataloader_dict, epsilons, attack_type, net_type, redetect_edge=False)
    print('Accuracy of robust redetect  model on adversarial images: %f %%' % acc_attack[0])
    fo.write('Accuracy of robust redetect  model on adversarial images: %f \n' % acc_attack[0])


    acc_attack, images = test_model_attack(net_robust, dataloader_dict, epsilons, attack_type, net_type, redetect_edge=True)
    print('Accuracy of robust redtect model on adversarial images with redetect_edge: %f %%' % acc_attack[0])
    fo.write('Accuracy of robust redetect model on adversarial images with redetect_edge: %f \n' % acc_attack[0])


fo.close()