def test(test_triplets, model, test_graph, all_triplets):
entity_embedding = model(test_graph.entity, test_graph.edge_index,
test_graph.edge_type, test_graph.edge_norm)
mrr = calc_mrr(entity_embedding,
model.relation_embedding,
test_triplets,
all_triplets,
hits=[1, 3, 10])
return mrr
def main(args):
# load graph data
data = load_data(args.dataset)
num_nodes = data.num_nodes
train_data = data.train
valid_data = data.valid
test_data = data.test
num_rels = data.num_rels
# check cuda
use_cuda = args.gpu >= 0 and torch.cuda.is_available()
if use_cuda:
torch.cuda.set_device(args.gpu)
# create model
model = LinkPredict(
model_class=RGCN, #KGVAE, #RGCN
in_dim=num_nodes,
h_dim=args.n_hidden,
num_rels=num_rels,
num_bases=args.n_bases,
num_hidden_layers=args.n_layers,
dropout=args.dropout,
use_cuda=True,
reg_param=args.regularization,
kl_param=args.kl_param,
)
# validation and testing triplets
valid_data = torch.LongTensor(valid_data)
test_data = torch.LongTensor(test_data)
# build val graph
val_adj_list, val_degrees = utils.get_adj_and_degrees(
num_nodes, valid_data)
if use_cuda:
model.cuda()
# build adj list and calculate degrees for sampling
adj_list, degrees = utils.get_adj_and_degrees(num_nodes, train_data)
# optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
model_state_file = 'model_state.pth'
forward_time = []
backward_time = []
# training loop
print("start training...")
epoch = 0
best_mrr = 0
train_records = open("train_records.txt", "w")
val_records = open("val_records.txt", "w")
while True:
model.train()
epoch += 1
# perform edge neighborhood sampling to generate training graph and data
g, node_id, edge_type, node_norm, data, labels = \
utils.generate_sampled_graph_and_labels(
train_data, args.graph_batch_size, args.graph_split_size,
num_rels, adj_list, degrees, args.negative_sample,
args.edge_sampler)
# print("Done edge sampling")
# set node/edge feature
node_id = torch.from_numpy(node_id).view(-1, 1).long()
edge_type = torch.from_numpy(edge_type)
edge_norm = node_norm_to_edge_norm(
g,
torch.from_numpy(node_norm).view(-1, 1))
data, labels = torch.from_numpy(data), torch.from_numpy(labels)
deg = g.in_degrees(range(g.number_of_nodes())).float().view(-1, 1)
if use_cuda:
node_id, deg = node_id.cuda(), deg.cuda()
edge_type, edge_norm = edge_type.cuda(), edge_norm.cuda()
data, labels = data.cuda(), labels.cuda()
t0 = time.time()
embed = model(g, node_id, edge_type, edge_norm)
loss, pred_loss, kl = model.get_loss(g, embed, data, labels)
train_records.write("{:d};{:.4f};{:.4f};{:.4f}\n".format(
epoch, loss.item(), pred_loss.item(), kl.item()))
train_records.flush()
t1 = time.time()
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(),
args.grad_norm) # clip gradients
optimizer.step()
t2 = time.time()
forward_time.append(t1 - t0)
backward_time.append(t2 - t1)
print(
"Epoch {:04d} | Loss {:.4f} | Best MRR {:.4f} | pred {:.4f} | reg {:.4f} | kl {:.4f} "
.format(epoch, loss.item(), best_mrr, pred_loss,
loss - pred_loss - kl, kl))
optimizer.zero_grad()
# validation
if epoch % args.evaluate_every == 1:
val_g, val_node_id, val_edge_type, val_node_norm, val_data, val_labels = utils.generate_sampled_graph_and_labels(
valid_data, args.graph_batch_size, args.graph_split_size,
num_rels, val_adj_list, val_degrees, args.negative_sample,
args.edge_sampler)
# print("Done edge sampling for validation")
val_node_id = torch.from_numpy(val_node_id).view(-1, 1).long()
val_edge_type = torch.from_numpy(val_edge_type)
val_edge_norm = node_norm_to_edge_norm(
val_g,
torch.from_numpy(val_node_norm).view(-1, 1))
val_data, val_labels = torch.from_numpy(
val_data), torch.from_numpy(val_labels)
if use_cuda:
val_node_id = val_node_id.cuda()
val_edge_type, val_edge_norm = val_edge_type.cuda(
), val_edge_norm.cuda()
val_data, val_neg_samples = val_data.cuda(), val_labels.cuda()
embed = model(val_g, val_node_id, val_edge_type, val_edge_norm)
mr, mrr, hits = utils.calc_mrr(embed,
model.w_relation,
val_data[val_labels == 1],
hits=[1, 3, 10],
eval_bz=args.eval_batch_size)
val_records.write("{:d};{:.4f};{:.4f};\n".format(epoch, mr, mrr) +
';'.join([str(i) for i in hits]) + "\n")
val_records.flush()
if mrr < best_mrr:
if epoch >= args.n_epochs:
break
else:
best_mrr = mrr
print("Best mrr", mrr)
torch.save({
'state_dict': model.state_dict(),
'epoch': epoch
}, model_state_file)
print("training done")
print("Mean forward time: {:4f}s".format(np.mean(forward_time)))
print("Mean Backward time: {:4f}s".format(np.mean(backward_time)))
def main(args):
# load graph data
data = load_data(args.dataset)
num_nodes = data.num_nodes
train_data = data.train
valid_data = data.valid
test_data = data.test
num_rels = data.num_rels
# check cuda
use_cuda = args.gpu >= 0 and torch.cuda.is_available()
if use_cuda:
torch.cuda.set_device(args.gpu)
# create model
model = LinkPredict(num_nodes,
args.n_hidden,
num_rels,
num_bases=args.n_bases,
num_hidden_layers=args.n_layers,
dropout=args.dropout,
use_cuda=use_cuda,
reg_param=args.regularization)
# validation and testing triplets
valid_data = torch.LongTensor(valid_data)
test_data = torch.LongTensor(test_data)
# build test graph
test_graph, test_rel, test_norm = utils.build_test_graph(
num_nodes, num_rels, train_data)
test_deg = test_graph.in_degrees(
range(test_graph.number_of_nodes())).float().view(-1,1)
test_node_id = torch.arange(0, num_nodes, dtype=torch.long).view(-1, 1)
test_rel = torch.from_numpy(test_rel)
test_norm = node_norm_to_edge_norm(test_graph, torch.from_numpy(test_norm).view(-1, 1))
if use_cuda:
model.cuda()
# build adj list and calculate degrees for sampling
adj_list, degrees = utils.get_adj_and_degrees(num_nodes, train_data)
# optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
model_state_file = 'model_state.pth'
forward_time = []
backward_time = []
# training loop
print("start training...")
epoch = 0
best_mrr = 0
while True:
model.train()
epoch += 1
# perform edge neighborhood sampling to generate training graph and data
g, node_id, edge_type, node_norm, data, labels = \
utils.generate_sampled_graph_and_labels(
train_data, args.graph_batch_size, args.graph_split_size,
num_rels, adj_list, degrees, args.negative_sample,
args.edge_sampler)
print("Done edge sampling")
# set node/edge feature
node_id = torch.from_numpy(node_id).view(-1, 1).long()
edge_type = torch.from_numpy(edge_type)
edge_norm = node_norm_to_edge_norm(g, torch.from_numpy(node_norm).view(-1, 1))
data, labels = torch.from_numpy(data), torch.from_numpy(labels)
deg = g.in_degrees(range(g.number_of_nodes())).float().view(-1, 1)
if use_cuda:
node_id, deg = node_id.cuda(), deg.cuda()
edge_type, edge_norm = edge_type.cuda(), edge_norm.cuda()
data, labels = data.cuda(), labels.cuda()
g = g.to(args.gpu)
t0 = time.time()
embed = model(g, node_id, edge_type, edge_norm)
loss = model.get_loss(g, embed, data, labels)
t1 = time.time()
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_norm) # clip gradients
optimizer.step()
t2 = time.time()
forward_time.append(t1 - t0)
backward_time.append(t2 - t1)
print("Epoch {:04d} | Loss {:.4f} | Best MRR {:.4f} | Forward {:.4f}s | Backward {:.4f}s".
format(epoch, loss.item(), best_mrr, forward_time[-1], backward_time[-1]))
optimizer.zero_grad()
# validation
if epoch % args.evaluate_every == 0:
# perform validation on CPU because full graph is too large
if use_cuda:
model.cpu()
model.eval()
print("start eval")
embed = model(test_graph, test_node_id, test_rel, test_norm)
mrr = utils.calc_mrr(embed, model.w_relation, torch.LongTensor(train_data),
valid_data, test_data, hits=[1, 3, 10], eval_bz=args.eval_batch_size,
eval_p=args.eval_protocol)
# save best model
if mrr < best_mrr:
if epoch >= args.n_epochs:
break
else:
best_mrr = mrr
torch.save({'state_dict': model.state_dict(), 'epoch': epoch},
model_state_file)
if use_cuda:
model.cuda()
print("training done")
print("Mean forward time: {:4f}s".format(np.mean(forward_time)))
print("Mean Backward time: {:4f}s".format(np.mean(backward_time)))
print("\nstart testing:")
# use best model checkpoint
checkpoint = torch.load(model_state_file)
if use_cuda:
model.cpu() # test on CPU
model.eval()
model.load_state_dict(checkpoint['state_dict'])
print("Using best epoch: {}".format(checkpoint['epoch']))
embed = model(test_graph, test_node_id, test_rel, test_norm)
utils.calc_mrr(embed, model.w_relation, torch.LongTensor(train_data), valid_data,
test_data, hits=[1, 3, 10], eval_bz=args.eval_batch_size, eval_p=args.eval_protocol)