-
Notifications
You must be signed in to change notification settings - Fork 1
/
test_mnist.py
82 lines (68 loc) · 3.11 KB
/
test_mnist.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 torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
from functions import test, adversary_test_attacks, test_adversary
from utils.utils import args_for_test_cifar
from models.cnn_mnist import ResNet_MNIST
import torch
import os
import csv
# ............. Get Basics ..................
# Get all the arguments
args = args_for_test_cifar()
state = {k: v for k, v in args._get_kwargs()}
# Use CUDA
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id
use_cuda = torch.cuda.is_available()
# ............. Load Data ...................
print('==> Preparing dataset %s' % args.dataset)
from utils.data import select_dataset
trainloader, testloader, num_classes = select_dataset(args)
# ......... create and load main model .........
title = '{}-{}-{}-{}'.format(args.dataset,args.model, args.depth, args.norm)
print("==> creating model: ", title)
model = ResNet_MNIST(args.depth, norm_type=args.norm)
print(' Total params: %.4fM' % (sum(p.numel() for p in model.parameters()) / 1000000.0))
# set model and its secondaries
model = model.cuda()
model = torch.nn.DataParallel(model)
cudnn.benchmark = True
criterion = nn.CrossEntropyLoss() # loss function
# Load checkpoint.
print('==> Loading checkpoint')
assert os.path.isfile(args.resume), 'Error: no checkpoint directory found!'
args.checkpoint = os.path.dirname(args.resume)
checkpoint = torch.load(args.resume)
model.load_state_dict(checkpoint['state_dict'])
# ............ find results ................
epsilons = args.epsilons
targeted = args.targeted
resume_path = args.resume
file_path = args.output_file_path
loss_normal, top1_acc_normal, _ = test(testloader, model, criterion, use_cuda, test_flag=False)
for epsilon in epsilons:
cw_initial_const = 10
attacks = ["FGSM", "LinfBIM", "MIM", "CWL2", "LinfPGD"]
attacks_loss, attacks_acc, _ = adversary_test_attacks(model, criterion, testloader, num_classes, attacks,
epsilon=epsilon, targeted=targeted, cw_initial_const=cw_initial_const,
use_cuda=True, test_flag=False)
# ........... save results to csv ................
csv_column_names = ["Model", "Targeted", "Training", "Dataset", "Norm", "Epsilon", "Normal", "FGSM", "LinfBIM", "MIM", "CWL2",
"LinfPGD"]
# if file does not exists, create it with header
if not os.path.exists(file_path):
print("path not exists!")
with open(file_path, 'w', newline='') as outcsv:
writer = csv.writer(outcsv)
writer.writerow(csv_column_names)
type = resume_path.split("/")[-2].split("-")[-1]
attack_names = ["Normal", "FGSM", "LinfBIM", "MIM", "CWL2", "LinfPGD"]
# convert dictionary to list to be fed to csv
acc = [args.model+str(args.depth), args.targeted, type, args.dataset, args.norm, epsilon, top1_acc_normal]
for attack_name in attacks:
acc.append(attacks_acc[attack_name])
assert(len(csv_column_names) == len(acc))
# save it in csv
with open(file_path, 'a', newline='') as outcsv:
writer = csv.writer(outcsv)
writer.writerow(acc)