-
Notifications
You must be signed in to change notification settings - Fork 0
/
test.py
66 lines (55 loc) · 2.45 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
CUDA_VISIBLE_DEVICES="0"
from validation import epochVal_metrics_test
import numpy as np
import torch.backends.cudnn as cudnn
import random
from cifar_load import get_dataloader, partition_data_allnoniid, record_net_data_stats
import copy
import torch
import torch.optim
import torch.nn.functional as F
from options import args_parser
from networks.models import ModelFedCon
#import matplotlib.pyplot as plt
#from mpl_toolkits.mplot3d import Axes3D
# import json
args = args_parser()
save_mode_path='final_model/SVHN_best.pth'
if __name__ == "__main__":
cudnn.benchmark = False
cudnn.deterministic = True
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
partition = torch.load('partition_strategy/SVHN_noniid_10%labeled.pth')
net_dataidx_map = partition['data_partition']
X_train, y_train, X_test, y_test, _, traindata_cls_counts = partition_data_allnoniid(
args.dataset, args.datadir, partition=args.partition, n_parties=args.num_users, beta=args.beta)
if args.dataset == 'SVHN':
X_train = X_train.transpose([0, 2, 3, 1])
X_test = X_test.transpose([0, 2, 3, 1])
if args.dataset == 'SVHN':
n_classes = 10
elif args.dataset == 'cifar100':
n_classes = 100
elif args.dataset == 'skin':
n_classes = 7
all_client = np.zeros([n_classes, n_classes], dtype=int)
traindata_cls_counts = record_net_data_stats(y_train, net_dataidx_map)
checkpoint_path = save_mode_path
checkpoint = torch.load(checkpoint_path)
net = ModelFedCon(args.model, args.out_dim, n_classes=n_classes)
model = net.cuda()
model.load_state_dict(checkpoint['state_dict'])
if args.dataset == 'SVHN' or args.dataset =='cifar100':
test_dl, test_ds = get_dataloader(args, X_test, y_test,
args.dataset, args.datadir, args.batch_size,
is_labeled=True, is_testing=True)
elif args.dataset == 'skin':
test_dl, test_ds = get_dataloader(args, X_test, y_test,
args.dataset, args.datadir, args.batch_size,
is_labeled=True, is_testing=True,pre_sz = args.pre_sz,input_sz = args.input_sz)
AUROCs, Accus, Pre, Recall = epochVal_metrics_test(model, test_dl,args.model, n_classes=n_classes)
AUROC_avg = np.array(AUROCs).mean()
Accus_avg = np.array(Accus).mean()