-
Notifications
You must be signed in to change notification settings - Fork 2
/
gaan.py
147 lines (126 loc) · 5.02 KB
/
gaan.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
from model_gaan import Model
from utils import *
from sklearn.metrics import roc_auc_score
import random
import dgl
from sklearn.metrics import average_precision_score
import argparse
from tqdm import tqdm
# os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(map(str, [0]))
# os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
# Set argument
parser = argparse.ArgumentParser(description='')
parser.add_argument('--dataset', type=str,
default='Amazon_no_isolate') # 'BlogCatalog' 'Flickr' 'ACM' 'cora' 'citeseer' 'pubmed'
parser.add_argument('--lr', type=float)
parser.add_argument('--weight_decay', type=float, default=0.0)
parser.add_argument('--seed', type=int, default=0)
parser.add_argument('--embedding_dim', type=int, default=300)
parser.add_argument('--num_epoch', type=int)
parser.add_argument('--drop_prob', type=float, default=0.0)
parser.add_argument('--batch_size', type=int, default=300)
parser.add_argument('--subgraph_size', type=int, default=4)
parser.add_argument('--readout', type=str, default='avg') # max min avg weighted_sum
parser.add_argument('--auc_test_rounds', type=int, default=256)
parser.add_argument('--negsamp_ratio', type=int, default=1)
args = parser.parse_args()
if args.lr is None:
if args.dataset in ['Amazon']:
args.lr = 1e-3
elif args.dataset in ['tf_finace']:
args.lr = 5e-4
elif args.dataset in ['reddit']:
args.lr = 1e-3
elif args.dataset in ['photo']:
args.lr = 1e-3
elif args.dataset in ['elliptic']:
args.lr = 5e-3
if args.num_epoch is None:
if args.dataset in ['reddit']:
args.num_epoch = 500
elif args.dataset in ['tf_finace']:
args.num_epoch = 1500
elif args.dataset in ['Amazon']:
args.num_epoch = 800
elif args.dataset in ['photo']:
args.num_epoch = 300
elif args.dataset in ['elliptic']:
args.num_epoch = 600
batch_size = args.batch_size
subgraph_size = args.subgraph_size
print('Dataset: ', args.dataset)
# Set random seed
dgl.random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
# torch.cuda.manual_seed(args.seed)
# torch.cuda.manual_seed_all(args.seed)
random.seed(args.seed)
# Load and preprocess data
adj, features, labels, all_idx, idx_train, idx_val, \
idx_test, ano_label, str_ano_label, attr_ano_label, normal_label_idx, abnormal_label_idx = load_mat(args.dataset)
if args.dataset in ['Amazon', 'tf_finace', 'reddit', 'elliptic']:
features, _ = preprocess_features(features)
else:
features = features.todense()
dgl_graph = adj_to_dgl_graph(adj)
nb_nodes = features.shape[0]
ft_size = features.shape[1]
# nb_classes = labels.shape[1]
raw_adj = adj
adj = normalize_adj(adj)
adj = (adj + sp.eye(adj.shape[0])).todense()
raw_adj = (raw_adj + sp.eye(raw_adj.shape[0])).todense()
features = torch.FloatTensor(features[np.newaxis])
adj = torch.FloatTensor(adj[np.newaxis])
raw_adj = torch.FloatTensor(raw_adj[np.newaxis])
labels = torch.FloatTensor(labels[np.newaxis])
# idx_train = torch.LongTensor(idx_train)
# idx_val = torch.LongTensor(idx_val)
# idx_test = torch.LongTensor(idx_test)
# Initialize model and optimiser
model = Model(ft_size, args.embedding_dim, 'prelu', args.negsamp_ratio, args.readout)
optimiser = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
optimiser_gen = torch.optim.Adam(model.generator.parameters(),
lr=args.lr)
# if torch.cuda.is_available():
# print('Using CUDA')
# model.cuda()
# features = features.cuda()
# adj = adj.cuda()
# labels = labels.cuda()
# idx_train = idx_train.cuda()
# idx_val = idx_val.cuda()
# idx_test = idx_test.cuda()
cnt_wait = 0
best = 1e9
best_t = 0
batch_num = nb_nodes // batch_size + 1
import time
# Train model
with tqdm(total=args.num_epoch) as pbar:
pbar.set_description('Training')
total_time = 0
for epoch in range(args.num_epoch):
start_time = time.time()
model.train()
optimiser.zero_grad()
optimiser_gen.zero_grad()
# Train model
# loss, loss_g, score_test = model(features, adj, normal_label_idx, idx_test)
loss, loss_g, score_test = model(features, adj, all_idx, idx_test)
loss.backward()
loss_g.backward()
optimiser.step()
optimiser_gen.step()
score_test = np.array(score_test.detach().cpu())
if epoch % 5 == 0:
print("Epoch:", '%04d' % (epoch), "train_loss=", "{:.5f}".format(loss.item()))
model.eval()
auc = roc_auc_score(ano_label[idx_test], score_test)
print('Testing {} AUC:{:.4f}'.format(args.dataset, auc))
AP = average_precision_score(ano_label[idx_test], score_test, average='macro', pos_label=1, sample_weight=None)
print('Testing AP:', AP)
print('Total time is', total_time)
end_time = time.time()
total_time += end_time - start_time