-
Notifications
You must be signed in to change notification settings - Fork 0
/
utils.py
79 lines (75 loc) · 2.45 KB
/
utils.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 os
import torch
import random
import numpy as np
import math
from models import GCN, SGC, SGC_rich, SAGE, APPNP, GAT
import torch.nn.functional as F
from copy import deepcopy
#%%
def seed(seed=0):
random.seed(seed)
os.environ['PYTHONHASHSEED'] = str(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
def get_GNN(modelname):
if modelname=='GCN':
GNN = GCN
elif modelname=='GAT':
GNN = GAT
elif modelname=='APPNP':
GNN = APPNP
elif modelname=='SAGE':
GNN = SAGE
elif modelname=='SGC':
GNN = SGC
elif modelname=='SGC_rich':
GNN = SGC_rich
return GNN
def evalue_model(model, x_test, adj_t_test, y_test, loss_fn=F.nll_loss, critic='acc'):
model.eval()
test_mask = y_test>=0
logits = model(x_test, adj_t_test)[test_mask]
loss = loss_fn(logits, y_test[test_mask])
if critic=='acc':
result = (logits.argmax(1) == y_test[test_mask]).sum()/logits.shape[0]
else:
raise Exception(f'Unimplemented critic {critic}!!')
return loss.item(), result
def train_model(model, x_syn, adj_t_syn, y_syn, optimizer, epochs, loss_fn=F.nll_loss):
losses = []
for i in range(epochs):
optimizer.zero_grad()
logits = model(x_syn, adj_t_syn)
loss = loss_fn(logits, y_syn)
loss.backward()
optimizer.step()
losses.append(loss.item())
return losses
def train_model_ealystop(model, x_syn, adj_t_syn, y_syn, x_val, adj_t_val, y_val,
optimizer, epochs, loss_fn=F.nll_loss):
val_mask = y_val >= 0
best_acc = 0.
loss_train_list = []
for i in range(epochs):
model.train()
optimizer.zero_grad()
logits = model(x_syn, adj_t_syn)
loss = loss_fn(logits, y_syn)
loss.backward()
optimizer.step()
with torch.no_grad():
model.eval()
logits_val = model(x_val, adj_t_val)[val_mask]
acc = (logits_val.argmax(1) == y_val[val_mask]).sum()/logits_val.shape[0]
if acc > best_acc:
loss_train_list.append(loss.item())
best_acc = acc
# print(f'train loss:{loss_train.item():.4}, acc:{acc:.4}')
weights = deepcopy(model.state_dict())
model.load_state_dict(weights)
return loss_train_list