forked from fanghenshaometeor/Orthogonal-Multi-Path
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathtest.py
181 lines (151 loc) · 6.74 KB
/
test.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
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
176
177
178
179
180
181
"""
Created on Mon Mar 09 2020
@author: fanghenshao
"""
from __future__ import print_function
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.data as data
import torch.optim as optim
from torchvision import datasets, transforms
import os
import ast
import argparse
import numpy as np
import scipy.io as sio
from utils import setup_seed
# ======== fix data type ========
torch.set_default_tensor_type(torch.FloatTensor)
# ======== fix seed =============
setup_seed(666)
# ======== options ==============
parser = argparse.ArgumentParser(description='Test Deep Neural Networks')
# -------- file param. --------------
parser.add_argument('--data_dir',type=str,default='/media/Disk1/KunFang/data/CIFAR10/',help='file path for data')
parser.add_argument('--dataset',type=str,default='CIFAR10',help='data set name')
parser.add_argument('--arch',type=str,default='OMPc',help='architecture of OMP model, alternative value include OMPa, OMPb and OMPc')
parser.add_argument('--model',type=str,default='vgg16',help='model architecture name')
parser.add_argument('--model_path',type=str,default='./save/CIFAR10-VGG.pth',help='saved model path')
# -------- training param. ----------
parser.add_argument('--batch_size',type=int,default=256,help='batch size for training (default: 256)')
parser.add_argument('--gpu_id',type=str,default='0',help='gpu device index')
# -------- hyper parameters -------
parser.add_argument('--num_paths',type=int,default=10,help='number of orthogonal paths')
parser.add_argument('--num_classes',type=int,default=10,help='number of classes')
args = parser.parse_args()
# ======== GPU device ========
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_id
# -------- main function
def main():
# ======== data set preprocess =============
# ======== mean-variance normalization is removed
if args.dataset == 'CIFAR10':
transform = transforms.Compose([
transforms.ToTensor()
])
trainset = datasets.CIFAR10(root=args.data_dir, train=True, download=True, transform=transform)
testset = datasets.CIFAR10(root=args.data_dir, train=False, download=True, transform=transform)
else:
assert False, "Unknow dataset : {}".format(args.dataset)
trainloader = data.DataLoader(trainset, batch_size=args.batch_size, shuffle=False)
testloader = data.DataLoader(testset, batch_size=args.batch_size, shuffle=False)
train_num, test_num = len(trainset), len(testset)
print('-------- DATA INFOMATION --------')
print('---- dataset: '+args.dataset)
print('---- #train : %d'%train_num)
print('---- #test : %d'%test_num)
# ======== load network ========
checkpoint = torch.load(args.model_path, map_location=torch.device("cpu"))
if args.model == 'vgg11':
if args.arch == 'OMPc':
from model.OMP_c_vgg import vgg11_bn
net = vgg11_bn(args.num_classes, args.num_paths).cuda()
else:
assert False, "Unsupported {}+{}".format(args.arch, args.model)
elif args.model == 'vgg13':
if args.arch == 'OMPc':
from model.OMP_c_vgg import vgg13_bn
net = vgg13_bn(args.num_classes, args.num_paths).cuda()
else:
assert False, "Unsupported {}+{}".format(args.arch, args.model)
elif args.model == 'vgg16':
if args.arch == 'OMPa':
from model.OMP_a_vgg import vgg16_bn
net = vgg16_bn(args.num_classes, args.num_paths).cuda()
elif args.arch == 'OMPb':
from model.OMP_b_vgg import vgg16_bn
net = vgg16_bn(args.num_classes, args.num_paths).cuda()
elif args.arch == 'OMPc':
from model.OMP_c_vgg import vgg16_bn
net = vgg16_bn(args.num_classes, args.num_paths).cuda()
else:
assert False, "Unknown architecture : {}".format(args.arch)
elif args.model == 'vgg19':
if args.arch == 'OMPc':
from model.OMP_c_vgg import vgg19_bn
net = vgg19_bn(args.num_classes, args.num_paths).cuda()
else:
assert False, "Unsupported {}+{}".format(args.arch, args.model)
elif args.model == 'resnet20':
if args.arch == 'OMPa':
from model.OMP_a_resnet_v1 import resnet20
net = resnet20(args.num_classes, args.num_paths).cuda()
elif args.arch == 'OMPb':
from model.OMP_b_resnet_v1 import resnet20
net = resnet20(args.num_classes, args.num_paths).cuda()
elif args.arch == 'OMPc':
from model.OMP_c_resnet_v1 import resnet20
net = resnet20(args.num_classes, args.num_paths).cuda()
else:
assert False, "Unknown architecture : {}".format(args.arch)
elif args.model == 'resnet32':
if args.arch == 'OMPc':
from model.OMP_c_resnet_v1 import resnet32
net = resnet32(args.num_classes, args.num_paths).cuda()
else:
assert False, "Unsupported {}+{}".format(args.arch, args.model)
else:
assert False, "Unknown model : {}".format(args.model)
net.load_state_dict(checkpoint['state_dict'])
net.eval()
print('-------- MODEL INFORMATION --------')
print("---- arch : "+args.arch)
print('---- model: '+args.model)
print('---- saved path: '+args.model_path)
print('-------- START TESTING --------')
corr_tr, corr_te = evaluate(net, trainloader, testloader)
print('Train acc. of each path:')
print(' ', corr_tr/train_num)
print('Test acc. of each path:')
print(' ', corr_te/test_num)
return
# -------- test model ---------------
def evaluate(net, trainloader, testloader):
net.eval()
correct_train, correct_test = np.zeros(args.num_paths), np.zeros(args.num_paths)
with torch.no_grad():
# -------- compute the accs. of train, test set
for test in testloader:
images, labels = test
images, labels = images.cuda(), labels.cuda()
# ------- forward
_, all_logits = net(images, 'all')
for idx in range(args.num_paths):
logits = all_logits[idx]
logits = logits.detach()
_, pred = torch.max(logits.data, 1)
correct_test[idx] += (pred == labels).sum().item()
for train in trainloader:
images, labels = train
images, labels = images.cuda(), labels.cuda()
_, all_logits = net(images, 'all')
for idx in range(args.num_paths):
logits = all_logits[idx]
logits = logits.detach()
_, pred = torch.max(logits.data, 1)
correct_train[idx] += (pred == labels).sum().item()
return correct_train, correct_test
# -------- start point
if __name__ == '__main__':
main()