Spaces:
Runtime error
Runtime error
File size: 10,394 Bytes
d2a8669 |
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 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 |
import torch
import torch.nn.functional as F
from torch.autograd import Variable
import time
import numpy as np
from sklearn import metrics
from sklearn.metrics import f1_score
from tqdm import trange, tqdm
import torch.nn as nn
from CatGCN.layers import StackedGNN
from FairGNN.src.models.GCN import GCN
class ClusterGNNTrainer(object):
"""
Training a huge graph cluster partition strategy.
"""
def __init__(self, args, clustering_machine, neptune_run):
self.args = args
self.clustering_machine = clustering_machine
self.neptune_run = neptune_run
self.device = torch.device("cuda:{}".format(args.gpu) if torch.cuda.is_available() else "cpu")
print('device:', self.device)
self.class_weight = clustering_machine.class_weight.to(self.device)
self.create_model()
# mew part -- input?
#self.sens_model = GCN(95, 128, 0.5) # number of feat, number of hidden, dropout percentage
#self.adv_model = nn.Linear(128, 1)
# adversary optimizer
#self.optimizer_A = torch.optim.Adam(self.adv_model.parameters(), lr = args.lr, weight_decay = args.weight_decay) ## can add weight decay
#self.criterion = nn.BCEWithLogitsLoss()
#self.A_loss = 0
def create_model(self):
"""
Creating a StackedGNN and transferring to CPU/GPU.
"""
self.model = StackedGNN(self.args, self.clustering_machine.field_count, self.clustering_machine.field_size, self.clustering_machine.class_count)
self.model = self.model.to(self.device)
def generate_field_adjs(self, node_count):
# Normalization by P'' = Q^{-1/2}*P'*Q^{-1/2}, P' = P+probe*O.
field_adjs = torch.ones((node_count, self.clustering_machine.field_size, self.clustering_machine.field_size))
field_adjs += self.args.diag_probe * torch.eye(self.clustering_machine.field_size)
row_sum = self.clustering_machine.field_size + self.args.diag_probe
field_adjs = (1. / row_sum) * field_adjs
return field_adjs
def do_forward_pass(self, cluster):
"""
Making a forward pass with data from a given partition.
:param cluster: Cluster index.
:return average_loss: Average loss on the cluster.
:return node_count: Number of nodes.
"""
edges = self.clustering_machine.sg_edges[cluster].to(self.device)
macro_nodes = self.clustering_machine.sg_nodes[cluster].to(self.device)
train_nodes = self.clustering_machine.sg_train_nodes[cluster].to(self.device)
field_index = self.clustering_machine.sg_field_index[cluster].to(self.device)
field_adjs = self.generate_field_adjs(field_index.shape[0]).to(self.device)
target = self.clustering_machine.sg_targets[cluster].to(self.device).squeeze()
prediction = self.model(edges, field_index, field_adjs)
# todo add forward pass for estimator and adversary
average_loss = F.nll_loss(prediction[train_nodes], target[train_nodes], self.class_weight)
node_count = train_nodes.shape[0]
return average_loss, node_count
def do_validation(self, cluster, epoch):
"""
Making a validation with data from a given partition.
:param cluster: Cluster index.
:return average_loss: Average loss on the cluster.
:return node_count: Number of nodes.
"""
edges = self.clustering_machine.sg_edges[cluster].to(self.device)
macro_nodes = self.clustering_machine.sg_nodes[cluster].to(self.device)
val_nodes = self.clustering_machine.sg_val_nodes[cluster].to(self.device)
field_index = self.clustering_machine.sg_field_index[cluster].to(self.device)
field_adjs = self.generate_field_adjs(field_index.shape[0]).to(self.device)
target = self.clustering_machine.sg_targets[cluster].to(self.device).squeeze()
prediction = self.model(edges, field_index, field_adjs)
average_loss = F.nll_loss(prediction[val_nodes], target[val_nodes], self.class_weight)
node_count = val_nodes.shape[0]
return average_loss, node_count
def do_prediction(self, cluster):
"""
Scoring a cluster.
:param cluster: Cluster index.
:return average_loss: Average loss on the cluster.
:return node_count: Number of nodes.
:return prediction: Prediction matrix with probabilities.
:return target: Target vector.
"""
edges = self.clustering_machine.sg_edges[cluster].to(self.device)
macro_nodes = self.clustering_machine.sg_nodes[cluster].to(self.device)
test_nodes = self.clustering_machine.sg_test_nodes[cluster].to(self.device)
field_index = self.clustering_machine.sg_field_index[cluster].to(self.device)
field_adjs = self.generate_field_adjs(field_index.shape[0]).to(self.device)
target = self.clustering_machine.sg_targets[cluster].to(self.device).squeeze()
prediction = self.model(edges, field_index, field_adjs)
average_loss = F.nll_loss(prediction[test_nodes], target[test_nodes], self.class_weight)
node_count = test_nodes.shape[0]
target = target[test_nodes]
prediction = prediction[test_nodes,:]
return average_loss, node_count, prediction, target
def update_average_loss(self, batch_average_loss, node_count):
"""
Updating the average loss in the epoch.
:param batch_average_loss: Loss of the cluster.
:param node_count: Number of nodes in currently processed cluster.
:return average_loss: Average loss in the epoch.
"""
self.accumulated_loss = self.accumulated_loss + batch_average_loss.item()*node_count
self.node_count_seen = self.node_count_seen + node_count
average_loss = self.accumulated_loss / self.node_count_seen
return average_loss
def train_val_test(self):
"""
Training, validation, and test a model per epoch.
"""
print("Training, validation, and test started.\n")
train_start_time = time.perf_counter()
bad_counter = 0
best_loss = np.inf
best_epoch = 0
self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.args.lr, weight_decay=self.args.weight_decay)
# test for epochs
print(self.args.epochs)
for epoch in range(1, self.args.epochs+1):
epoch_start_time = time.time()
np.random.shuffle(self.clustering_machine.clusters)
self.model.train()
self.node_count_seen = 0
self.accumulated_loss = 0
for cluster in self.clustering_machine.clusters:
self.optimizer.zero_grad()
batch_average_loss, node_count = self.do_forward_pass(cluster)
batch_average_loss.backward()
self.optimizer.step()
average_loss = self.update_average_loss(batch_average_loss, node_count)
train_loss = average_loss
self.model.eval()
self.node_count_seen = 0
self.accumulated_loss = 0
for cluster in self.clustering_machine.clusters:
batch_average_loss, node_count = self.do_validation(cluster, epoch)
average_loss = self.update_average_loss(batch_average_loss, node_count)
val_loss = average_loss
print("Epoch: {:04d}".format(epoch),
"||",
"time cost: {:.2f}s".format(time.time() - epoch_start_time),
"||",
"train loss: {:.4f}".format(train_loss),
"val loss: {:.4f}".format(val_loss))
if val_loss < best_loss:
best_loss = val_loss
best_epoch = epoch
bad_counter = 0
best_model_state = self.model.state_dict()
else:
bad_counter += 1
if bad_counter == self.args.patience:
break
self.model.load_state_dict(best_model_state)
self.model.eval()
self.node_count_seen = 0
self.accumulated_loss = 0
self.predictions = []
self.targets = []
for cluster in self.clustering_machine.clusters:
batch_average_loss, node_count, prediction, target = self.do_prediction(cluster)
average_loss = self.update_average_loss(batch_average_loss, node_count)
self.predictions.append(prediction.cpu().detach().numpy())
self.targets.append(target.cpu().detach().numpy())
test_loss = average_loss
self.targets = np.concatenate(self.targets)
self.predictions = np.concatenate(self.predictions).argmax(1)
acc_score = metrics.accuracy_score(self.targets, self.predictions)
macro_f1 = metrics.f1_score(self.targets, self.predictions, average="macro")
classification_report = metrics.classification_report(self.targets, self.predictions, digits=4)
print(classification_report)
# Confusion matrics and AUC
confusion_matrix = metrics.confusion_matrix(self.targets, self.predictions)
print(confusion_matrix)
#F1
f1 = f1_score(self.targets, self.predictions, average='macro')
print('F1 score:', f1)
# fpr, tpr, _ = metrics.roc_curve(self.targets, self.predictions)
# auc = metrics.auc(fpr, tpr)
# print("AUC:", auc)
train_time = (time.perf_counter() - train_start_time)/60
print("Optimization Finished!")
print("Total time elapsed: {:.2f}min".format(train_time))
print("Best Result:\n",
"best epoch: {:04d}".format(best_epoch),
"||",
"test loss: {:.4f}".format(test_loss),
"||",
"accuracy: {:.4f}".format(acc_score),
"macro-f1: {:.4f}".format(macro_f1))
# Save results on Neptune
self.neptune_run["best_epoch"] = best_epoch
self.neptune_run["test/loss"] = test_loss
self.neptune_run["test/acc"] = acc_score
self.neptune_run["test/f1"] = macro_f1
# self.neptune_run["test/auc"] = auc
# self.neptune_run["test/tpr"] = tpr
# self.neptune_run["test/fpr"] = fpr
self.neptune_run["conf_matrix"] = confusion_matrix
self.neptune_run["train_time"] = train_time
|