-
Notifications
You must be signed in to change notification settings - Fork 3
/
Copy pathtest4tune.py
141 lines (114 loc) · 4.36 KB
/
test4tune.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
import os
import os.path as osp
import sys
import glob
import numpy as np
import torch
import utils
import logging
import pickle
import argparse
import torch.nn as nn
import genotypes
import torch.utils
import torchvision.datasets as dset
import torch.backends.cudnn as cudnn
import torch.nn.functional as F
from torch import cat
from torch.autograd import Variable
from model import NetworkGNN as Network
from torch_geometric.datasets import Planetoid, Amazon, Coauthor, CoraFull, Reddit
from sklearn.model_selection import StratifiedKFold
from torch_geometric.utils import add_self_loops
from logging_util import init_logger
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
def index_to_mask(index, size):
mask = torch.zeros(size, dtype=torch.bool, device=index.device)
mask[index] = 1
return mask
def gen_uniform_60_20_20_split(data):
skf = StratifiedKFold(5, shuffle=True)
idx = [torch.from_numpy(i) for _, i in skf.split(data.y, data.y)]
return cat(idx[:3], 0), cat(idx[3:4], 0), cat(idx[4:], 0)
def save_load_split(data, raw_dir, run, gen_splits):
prefix = gen_splits.__name__[4:-6]
path = osp.join(raw_dir, '..', '{}_{:03d}.pt'.format(prefix, run))
if osp.exists(path):
split = torch.load(path)
else:
split = gen_splits(data)
torch.save(split, path)
data.train_mask = index_to_mask(split[0], data.num_nodes)
data.val_mask = index_to_mask(split[1], data.num_nodes)
data.test_mask = index_to_mask(split[2], data.num_nodes)
return data
def main(test_args1):
global test_args
test_args = test_args1
if not torch.cuda.is_available():
logging.info('no gpu device available')
sys.exit(1)
#np.random.seed(test_args.seed)
torch.cuda.set_device(test_args.gpu)
cudnn.benchmark = True
torch.manual_seed(test_args.seed)
cudnn.enabled=True
torch.cuda.manual_seed(test_args.seed)
#path = osp.join('../data', 'Cora')
if test_args.data == 'Amazon_Computers':
dataset = Amazon('../data/Amazon_Computers', 'Computers')
elif test_args.data == 'Coauthor_Physics':
dataset = Coauthor('../data/Coauthor_Physics', 'Physics')
elif test_args.data == 'Coauthor_CS':
dataset = Coauthor('../data/Coauthor_CS', 'CS')
elif test_args.data == 'Cora_Full':
dataset = CoraFull('../data/Cora_Full')
elif test_args.data == 'PubMed':
dataset = Planetoid('../data/PubMed', 'PubMed')
elif test_args.data == 'Cora':
dataset = Planetoid('../data/Cora', 'Cora')
elif test_args.data == 'CiteSeer':
dataset = Planetoid('../data/CiteSeer', 'CiteSeer')
if test_args.data == 'small_Reddit':
dataset = Reddit('../data/Reddit/')
with open('../data/small_Reddit/sampled_reddit.obj', 'rb') as f:
data = pickle.load(f)
raw_dir = '../data/small_Reddit/raw/'
else:
raw_dir = dataset.raw_dir
data = dataset[0]
data = save_load_split(data, raw_dir, test_args.rnd_num, gen_uniform_60_20_20_split)
edge_index, _ = add_self_loops(data.edge_index)
data.edge_index = edge_index
hidden_size = test_args.hidden_size
genotype = test_args.arch
criterion = nn.CrossEntropyLoss()
criterion = criterion.cuda()
model = Network(genotype, criterion, dataset.num_features, dataset.num_classes, hidden_size, num_layers=test_args.num_layers, in_dropout=test_args.in_dropout, out_dropout=test_args.out_dropout, act=test_args.activation)
model = model.cuda()
utils.load(model, test_args.model_path)
logging.info("gpu=%s, genotype=%s, param size = %fMB, args=%s", test_args.gpu, genotype, utils.count_parameters_in_MB(model), test_args.__dict__)
test_acc, test_obj = infer(data, model, criterion)
logging.info('test_acc=%f', test_acc)
return test_acc, test_args.save
def infer(data, model, criterion):
objs = utils.AvgrageMeter()
top1 = utils.AvgrageMeter()
top5 = utils.AvgrageMeter()
model.eval()
with torch.no_grad():
logits = F.log_softmax(model(data.to(device)), dim=-1)
input = logits[data.test_mask].to(device)
target = data.y[data.test_mask].to(device)
#logits, _ = model(input)
loss = criterion(input, target)
pred = logits[data.test_mask].max(1)[1]
acc = (pred == target).sum().item() / data.test_mask.sum().item()
prec1, prec5 = utils.accuracy(input, target, topk=(1, 3))
n = input.size(0)
objs.update(loss.data.item(), n)
top1.update(prec1.data.item(), n)
top5.update(prec5.data.item(), n)
return top1.avg, objs.avg
if __name__ == '__main__':
main()