def train(
dataset,
model,
args,
same_feat=True,
val_dataset=None,
test_dataset=None,
writer=None,
mask_nodes=True,
):
writer_batch_idx = [0, 3, 6, 9]
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad,
model.parameters()),
lr=0.001)
iter = 0
best_val_result = {"epoch": 0, "loss": 0, "acc": 0}
test_result = {"epoch": 0, "loss": 0, "acc": 0}
train_accs = []
train_epochs = []
best_val_accs = []
best_val_epochs = []
test_accs = []
test_epochs = []
val_accs = []
for epoch in range(args.num_epochs):
begin_time = time.time()
avg_loss = 0.0
model.train()
predictions = []
print("Epoch: ", epoch)
for batch_idx, data in enumerate(dataset):
model.zero_grad()
if batch_idx == 0:
prev_adjs = data["adj"]
prev_feats = data["feats"]
prev_labels = data["label"]
all_adjs = prev_adjs
all_feats = prev_feats
all_labels = prev_labels
elif batch_idx < 20:
prev_adjs = data["adj"]
prev_feats = data["feats"]
prev_labels = data["label"]
all_adjs = torch.cat((all_adjs, prev_adjs), dim=0)
all_feats = torch.cat((all_feats, prev_feats), dim=0)
all_labels = torch.cat((all_labels, prev_labels), dim=0)
adj = Variable(data["adj"].float(), requires_grad=False).cuda()
h0 = Variable(data["feats"].float(), requires_grad=False).cuda()
label = Variable(data["label"].long()).cuda()
batch_num_nodes = data["num_nodes"].int().numpy(
) if mask_nodes else None
assign_input = Variable(data["assign_feats"].float(),
requires_grad=False).cuda()
ypred, att_adj = model(h0,
adj,
batch_num_nodes,
assign_x=assign_input)
if batch_idx < 5:
predictions += ypred.cpu().detach().numpy().tolist()
if not args.method == "soft-assign" or not args.linkpred:
loss = model.loss(ypred, label)
else:
loss = model.loss(ypred, label, adj, batch_num_nodes)
loss.backward()
nn.utils.clip_grad_norm(model.parameters(), args.clip)
optimizer.step()
iter += 1
avg_loss += loss
avg_loss /= batch_idx + 1
elapsed = time.time() - begin_time
if writer is not None:
writer.add_scalar("loss/avg_loss", avg_loss, epoch)
if args.linkpred:
writer.add_scalar("loss/linkpred_loss", model.link_loss, epoch)
print("Avg loss: ", avg_loss, "; epoch time: ", elapsed)
result = evaluate(dataset,
model,
args,
name="Train",
max_num_examples=100)
train_accs.append(result["acc"])
train_epochs.append(epoch)
if val_dataset is not None:
val_result = evaluate(val_dataset, model, args, name="Validation")
val_accs.append(val_result["acc"])
if val_result["acc"] > best_val_result["acc"] - 1e-7:
best_val_result["acc"] = val_result["acc"]
best_val_result["epoch"] = epoch
best_val_result["loss"] = avg_loss
if test_dataset is not None:
test_result = evaluate(test_dataset, model, args, name="Test")
test_result["epoch"] = epoch
if writer is not None:
writer.add_scalar("acc/train_acc", result["acc"], epoch)
writer.add_scalar("acc/val_acc", val_result["acc"], epoch)
writer.add_scalar("loss/best_val_loss", best_val_result["loss"],
epoch)
if test_dataset is not None:
writer.add_scalar("acc/test_acc", test_result["acc"], epoch)
print("Best val result: ", best_val_result)
best_val_epochs.append(best_val_result["epoch"])
best_val_accs.append(best_val_result["acc"])
if test_dataset is not None:
print("Test result: ", test_result)
test_epochs.append(test_result["epoch"])
test_accs.append(test_result["acc"])
matplotlib.style.use("seaborn")
plt.switch_backend("agg")
plt.figure()
plt.plot(train_epochs,
math_utils.exp_moving_avg(train_accs, 0.85),
"-",
lw=1)
if test_dataset is not None:
plt.plot(best_val_epochs, best_val_accs, "bo", test_epochs, test_accs,
"go")
plt.legend(["train", "val", "test"])
else:
plt.plot(best_val_epochs, best_val_accs, "bo")
plt.legend(["train", "val"])
plt.savefig(io_utils.gen_train_plt_name(args), dpi=600)
plt.close()
matplotlib.style.use("default")
print(all_adjs.shape, all_feats.shape, all_labels.shape)
cg_data = {
"adj": all_adjs,
"feat": all_feats,
"label": all_labels,
"pred": np.expand_dims(predictions, axis=0),
"train_idx": list(range(len(dataset))),
}
io_utils.save_checkpoint(model,
optimizer,
args,
num_epochs=-1,
cg_dict=cg_data)
return model, val_accs
def train_node_classifier(G, labels, model, args, writer=None):
# train/test split only for nodes
num_nodes = G.number_of_nodes()
num_train = int(num_nodes * args.train_ratio)
idx = [i for i in range(num_nodes)]
np.random.shuffle(idx)
train_idx = idx[:num_train]
test_idx = idx[num_train:]
data = gengraph.preprocess_input_graph(G, labels)
labels_train = torch.tensor(data["labels"][:, train_idx], dtype=torch.long)
adj = torch.tensor(data["adj"], dtype=torch.float)
x = torch.tensor(data["feat"], requires_grad=True, dtype=torch.float)
scheduler, optimizer = train_utils.build_optimizer(
args, model.parameters(), weight_decay=args.weight_decay)
model.train()
ypred = None
for epoch in range(args.num_epochs):
begin_time = time.time()
model.zero_grad()
if args.gpu:
ypred, adj_att = model(x.cuda(), adj.cuda())
else:
ypred, adj_att = model(x, adj)
ypred_train = ypred[:, train_idx, :]
if args.gpu:
loss = model.loss(ypred_train, labels_train.cuda())
else:
loss = model.loss(ypred_train, labels_train)
loss.backward()
nn.utils.clip_grad_norm(model.parameters(), args.clip)
optimizer.step()
#for param_group in optimizer.param_groups:
# print(param_group["lr"])
elapsed = time.time() - begin_time
result_train, result_test = evaluate_node(ypred.cpu(), data["labels"],
train_idx, test_idx)
if writer is not None:
writer.add_scalar("loss/avg_loss", loss, epoch)
writer.add_scalars(
"prec",
{
"train": result_train["prec"],
"test": result_test["prec"]
},
epoch,
)
writer.add_scalars(
"recall",
{
"train": result_train["recall"],
"test": result_test["recall"]
},
epoch,
)
writer.add_scalars("acc", {
"train": result_train["acc"],
"test": result_test["acc"]
}, epoch)
if epoch % 10 == 0:
print(
"epoch: ",
epoch,
"; loss: ",
loss.item(),
"; train_acc: ",
result_train["acc"],
"; test_acc: ",
result_test["acc"],
"; train_prec: ",
result_train["prec"],
"; test_prec: ",
result_test["prec"],
"; epoch time: ",
"{0:0.2f}".format(elapsed),
)
if scheduler is not None:
scheduler.step()
print(result_train["conf_mat"])
print(result_test["conf_mat"])
# computation graph
model.eval()
if args.gpu:
ypred, _ = model(x.cuda(), adj.cuda())
else:
ypred, _ = model(x, adj)
cg_data = {
"adj": data["adj"],
"feat": data["feat"],
"label": data["labels"],
"pred": ypred.cpu().detach().numpy(),
"train_idx": train_idx,
}
# import pdb
# pdb.set_trace()
io_utils.save_checkpoint(model,
optimizer,
args,
num_epochs=-1,
cg_dict=cg_data)
def train_node_classifier(G, labels, model, args, writer=None):
# train/test split only for nodes
num_nodes = G.number_of_nodes()
# Training data with 80% ratio, labels_train.size()
num_train = int(num_nodes * args.train_ratio)
idx = [i for i in range(num_nodes)]
# Shuffle for training
np.random.shuffle(idx)
train_idx = idx[:num_train]
test_idx = idx[num_train:]
data = gengraph.preprocess_input_graph(G, labels)
labels_train = torch.tensor(data["labels"][:, train_idx], dtype=torch.long)
adj = torch.tensor(data["adj"], dtype=torch.float)
x = torch.tensor(data["feat"], requires_grad=True, dtype=torch.float)
# scheduler, optimizer = train_utils.build_optimizer(
# args, model.parameters(), weight_decay=args.weight_decay
# )
# list(testModel.parameters()) and list(filter_fn) to show contents
# train_utils.build_optimizer
filter_fn = filter(lambda p: p.requires_grad, model.parameters())
# args.opt == 'adam':
optimizer = optim.Adam(filter_fn, lr=args.lr, weight_decay=0.0)
scheduler = None
# Sets the module in training mode
model.train()
ypred = None
for epoch in range(args.num_epochs):
begin_time = time.time()
model.zero_grad()
if args.gpu:
ypred, adj_att = model(x.cuda(), adj.cuda())
else:
ypred, adj_att = model(x, adj)
ypred_train = ypred[:, train_idx, :]
if args.gpu:
loss = model.loss(ypred_train, labels_train.cuda())
else:
loss = model.loss(ypred_train, labels_train)
loss.backward()
nn.utils.clip_grad_norm_(model.parameters(), args.clip)
optimizer.step()
#for param_group in optimizer.param_groups:
# print(param_group["lr"])
elapsed = time.time() - begin_time
# Obtain with Confusion matrices for Train and Test results
result_train, result_test = evaluate_node(ypred.cpu(), data["labels"],
train_idx, test_idx)
if writer is not None:
writer.add_scalar("loss/avg_loss", loss, epoch)
writer.add_scalars(
"prec",
{
"train": result_train["prec"],
"test": result_test["prec"]
},
epoch,
)
writer.add_scalars(
"recall",
{
"train": result_train["recall"],
"test": result_test["recall"]
},
epoch,
)
writer.add_scalars("acc", {
"train": result_train["acc"],
"test": result_test["acc"]
}, epoch)
if epoch % 10 == 0:
print(
"epoch: ",
epoch,
"; loss: ",
loss.item(),
"; train_acc: ",
result_train["acc"],
"; test_acc: ",
result_test["acc"],
"; train_prec: ",
result_train["prec"],
"; test_prec: ",
result_test["prec"],
"; epoch time: ",
"{0:0.2f}".format(elapsed),
)
if scheduler is not None:
scheduler.step()
print("Confusion Matrix of train result :\n", result_train["conf_mat"])
print("Confusion Matrix of test result :\n", result_test["conf_mat"])
# Sets the module in evaluation mode for computational graph
model.eval()
if args.gpu:
ypred, _ = model(x.cuda(), adj.cuda())
else:
ypred, _ = model(x, adj)
cg_data = {
"adj": data["adj"],
"feat": data["feat"],
"label": data["labels"],
"pred": ypred.cpu().detach().numpy(),
"train_idx": train_idx,
}
print("Labels of the Computational graph :\n", cg_data['label'])
print("Prediction result of the Computational graph :\n", cg_data['pred'])
print("Train index of the Computational graph data :\n",
cg_data['train_idx'])
# import pdb
# pdb.set_trace()
io_utils.save_checkpoint(model,
optimizer,
args,
num_epochs=-1,
cg_dict=cg_data)
def train_node_classifier_multigraph(G_list, labels, model, args, writer=None):
train_idx_all, test_idx_all = [], []
# train/test split only for nodes
num_nodes = G_list[0].number_of_nodes()
num_train = int(num_nodes * args.train_ratio)
idx = [i for i in range(num_nodes)]
np.random.shuffle(idx)
train_idx = idx[:num_train]
train_idx_all.append(train_idx)
test_idx = idx[num_train:]
test_idx_all.append(test_idx)
data = gengraph.preprocess_input_graph(G_list[0], labels[0])
all_labels = data["labels"]
labels_train = torch.tensor(data["labels"][:, train_idx], dtype=torch.long)
adj = torch.tensor(data["adj"], dtype=torch.float)
x = torch.tensor(data["feat"], requires_grad=True, dtype=torch.float)
for i in range(1, len(G_list)):
np.random.shuffle(idx)
train_idx = idx[:num_train]
train_idx_all.append(train_idx)
test_idx = idx[num_train:]
test_idx_all.append(test_idx)
data = gengraph.preprocess_input_graph(G_list[i], labels[i])
all_labels = np.concatenate((all_labels, data["labels"]), axis=0)
labels_train = torch.cat(
[
labels_train,
torch.tensor(data["labels"][:, train_idx], dtype=torch.long),
],
dim=0,
)
adj = torch.cat([adj, torch.tensor(data["adj"], dtype=torch.float)])
x = torch.cat([
x,
torch.tensor(data["feat"], requires_grad=True, dtype=torch.float)
])
scheduler, optimizer = train_utils.build_optimizer(
args, model.parameters(), weight_decay=args.weight_decay)
model.train()
ypred = None
for epoch in range(args.num_epochs):
begin_time = time.time()
model.zero_grad()
if args.gpu:
ypred = model(x.cuda(), adj.cuda())
else:
ypred = model(x, adj)
# normal indexing
ypred_train = ypred[:, train_idx, :]
# in multigraph setting we can't directly access all dimensions so we need to gather all the training instances
all_train_idx = [item for sublist in train_idx_all for item in sublist]
ypred_train_cmp = torch.cat(
[ypred[i, train_idx_all[i], :] for i in range(10)],
dim=0).reshape(10, 146, 6)
if args.gpu:
loss = model.loss(ypred_train_cmp, labels_train.cuda())
else:
loss = model.loss(ypred_train_cmp, labels_train)
loss.backward()
nn.utils.clip_grad_norm(model.parameters(), args.clip)
optimizer.step()
#for param_group in optimizer.param_groups:
# print(param_group["lr"])
elapsed = time.time() - begin_time
result_train, result_test = evaluate_node(ypred.cpu(), all_labels,
train_idx_all, test_idx_all)
if writer is not None:
writer.add_scalar("loss/avg_loss", loss, epoch)
writer.add_scalars(
"prec",
{
"train": result_train["prec"],
"test": result_test["prec"]
},
epoch,
)
writer.add_scalars(
"recall",
{
"train": result_train["recall"],
"test": result_test["recall"]
},
epoch,
)
writer.add_scalars("acc", {
"train": result_train["acc"],
"test": result_test["acc"]
}, epoch)
print(
"epoch: ",
epoch,
"; loss: ",
loss.item(),
"; train_acc: ",
result_train["acc"],
"; test_acc: ",
result_test["acc"],
"; epoch time: ",
"{0:0.2f}".format(elapsed),
)
if scheduler is not None:
scheduler.step()
print(result_train["conf_mat"])
print(result_test["conf_mat"])
# computation graph
model.eval()
if args.gpu:
ypred = model(x.cuda(), adj.cuda())
else:
ypred = model(x, adj)
cg_data = {
"adj": adj.cpu().detach().numpy(),
"feat": x.cpu().detach().numpy(),
"label": all_labels,
"pred": ypred.cpu().detach().numpy(),
"train_idx": train_idx_all,
}
io_utils.save_checkpoint(model,
optimizer,
args,
num_epochs=-1,
cg_dict=cg_data)
def medic(args):
"""
Creating a simple Graph ConvNet using parameters of args (https://arxiv.org/abs/1609.02907)
"""
# Loading DataSet from /Pickles
global result_test, result_train
with open('Pickles/feats.pickle', 'rb') as handle:
feats = np.expand_dims(pickle.load(handle), axis=0)
with open('Pickles/age_adj.pickle', 'rb') as handle:
age_adj = pickle.load(handle)
with open('Pickles/preds.pickle', 'rb') as handle:
labels = np.expand_dims(pickle.load(handle), axis=0)
# initializing model variables
num_nodes = labels.shape[1]
num_train = int(num_nodes * 0.9)
num_classes = max(labels[0]) + 1
idx = [i for i in range(num_nodes)]
np.random.shuffle(idx)
train_idx = idx[:num_train]
test_idx = idx[num_train:]
labels = labels.astype(np.long)
age_adj = age_adj.astype(np.float)
feats = feats.astype(np.float)
age_adj = age_adj + np.eye(age_adj.shape[0])
d_hat_inv = np.linalg.inv(np.diag(age_adj.sum(axis=1)))**(1 / 2)
temp = np.matmul(d_hat_inv, age_adj)
age_adj = np.matmul(temp, d_hat_inv)
age_adj = np.expand_dims(age_adj, axis=0)
labels_train = torch.tensor(labels[:, train_idx], dtype=torch.long)
adj = torch.tensor(age_adj, dtype=torch.float)
x = torch.tensor(feats, dtype=torch.float, requires_grad=True)
# Creating a model which is used in https://github.com/RexYing/gnn-model-explainer
model = models.GcnEncoderNode(
args.input_dim,
args.hidden_dim,
args.output_dim,
num_classes,
args.num_gc_layers,
bn=args.bn,
args=args,
)
if args.gpu:
model = model.cuda()
scheduler, optimizer = build_optimizer(args,
model.parameters(),
weight_decay=args.weight_decay)
model.train()
to_save = (0, None) # used for saving best model
# training the model
for epoch in range(args.num_epochs):
begin_time = time.time()
model.zero_grad()
if args.gpu:
ypred, adj_att = model(x.cuda(), adj.cuda())
else:
ypred, adj_att = model(x, adj)
ypred_train = ypred[:, train_idx, :]
if args.gpu:
loss = model.loss(ypred_train, labels_train.cuda())
else:
loss = model.loss(ypred_train, labels_train)
loss.backward()
nn.utils.clip_grad_norm(model.parameters(), args.clip)
optimizer.step()
# for param_group in optimizer.param_groups:
# print(param_group["lr"])
elapsed = time.time() - begin_time
result_train, result_test = evaluate_node(ypred.cpu(), labels,
train_idx, test_idx)
if result_test["acc"] > to_save[0]:
to_save = (result_test["acc"], (model, optimizer, args))
if epoch % 10 == 0:
print(
"epoch: ",
epoch,
"; loss: ",
loss.item(),
"; train_acc: ",
result_train["acc"],
"; test_acc: ",
result_test["acc"],
"; train_prec: ",
result_train["prec"],
"; test_prec: ",
result_test["prec"],
"; epoch time: ",
"{0:0.2f}".format(elapsed),
)
if epoch % 100 == 0:
print(result_train["conf_mat"])
print(result_test["conf_mat"])
if scheduler is not None:
scheduler.step()
print(result_train["conf_mat"])
print(result_test["conf_mat"])
to_save[1][0].eval()
if args.gpu:
ypred, _ = to_save[1][0](x.cuda(), adj.cuda())
else:
ypred, _ = to_save[1][0](x, adj)
cg_data = {
"adj": age_adj,
"feat": feats,
"label": labels,
"pred": ypred.cpu().detach().numpy(),
"train_idx": train_idx,
}
# saving the model so that it can be restored for GNN explaining
print(
save_checkpoint(to_save[1][0],
to_save[1][1],
args,
num_epochs=-1,
cg_dict=cg_data))
return to_save[1][0], to_save[1][1], args, cg_data