def train(self):
self.model.train()
g, node_id, edge_type, node_norm, data, labels = \
generate_sampled_graph_and_labels(triplets=self.train_data,
sample_size=self.p.graph_batch_size,
split_size=self.p.graph_split_size,
num_rels=self.num_rels,
adj_list=self.adj_list,
degrees=self.in_deg,
negative_rate=self.p.negative_sample,
sampler=self.p.edge_sampler)
node_id = torch.from_numpy(node_id).view(-1, 1).long().to(self.device)
edge_type = torch.from_numpy(edge_type).to(self.device)
data, labels = torch.from_numpy(data).to(
self.device), torch.from_numpy(labels).to(self.device)
edge_norm = node_norm_2_edge_norm(
g,
torch.from_numpy(node_norm).view(-1, 1)).to(self.device)
# deg = g.in_degrees(range(g.number_of_nodes())).float().view(-1, 1).to(self.device)
output = self.model(g, node_id, edge_type, edge_norm)
loss = self.model.calc_loss(output, data, labels)
self.optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(self.model.parameters(),
self.p.grad_norm)
self.optimizer.step()
return loss.item()
def train(train_triplets, model, use_cuda, batch_size, split_size,
negative_sample, reg_ratio, num_entities, num_relations):
train_data = generate_sampled_graph_and_labels(train_triplets, batch_size,
split_size, num_entities,
num_relations,
negative_sample)
if use_cuda:
device = torch.device('cuda')
train_data.to(device)
entity_embedding = model(train_data.entity, train_data.edge_index,
train_data.edge_type, train_data.edge_norm)
loss = model.score_loss(
entity_embedding, train_data.samples,
train_data.labels) + reg_ratio * model.reg_loss(entity_embedding)
return loss
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)
def main(args):
set_seeds(args.seed)
# load graph data
if args.dataset == "FB15K-237":
dataset_cls = FB15kReader
data_dir = "data/FB15k-237/"
elif args.dataset == "atomic":
dataset_cls = AtomicTSVReader
data_dir = "data/atomic/"
elif args.dataset == "conceptnet":
dataset_cls = ConceptNetTSVReader
data_dir = "data/ConceptNet/"
else:
raise ValueError("Invalid option for dataset.")
# Store entity-wise dicts for filtered metrics
train_data, valid_data, test_data, valid_labels, test_labels, train_network = load_data(
args.dataset, dataset_cls, data_dir, args.sim_relations)
num_nodes = len(train_network.graph.nodes)
num_rels = len(train_network.graph.relations)
all_tuples = train_data.tolist() + valid_data.tolist() + test_data.tolist()
# for filtered ranking
all_e1_to_multi_e2, all_e2_to_multi_e1 = reader_utils.create_entity_dicts(
all_tuples, num_rels, args.sim_relations)
# for training
train_e1_to_multi_e2, train_e2_to_multi_e1 = reader_utils.create_entity_dicts(
train_data.tolist(), num_rels, args.sim_relations)
# the below dicts `include` sim relations
sim_train_e1_to_multi_e2, sim_train_e2_to_multi_e1 = reader_utils.create_entity_dicts(
train_data.tolist(), num_rels)
# check cuda
use_cuda = torch.cuda.is_available()
if use_cuda and not args.no_cuda:
torch.cuda.set_device(args.gpu)
cpu_decoding = args.cpu_decoding
# atomic graph is much larger, so we perform evaluation on cpu
cpu_eval = True if args.dataset == "atomic" else False
# create model
model = LinkPredictor(num_nodes, num_rels, args, use_cuda=use_cuda)
# build graph
graph_train_data = train_data
test_graph, test_rel, test_norm = utils.build_test_graph(
num_nodes, num_rels, graph_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).view(-1, 1)
test_norm = torch.from_numpy(test_norm).view(-1, 1)
# transfer graph data to gpu
if use_cuda and not args.no_cuda and not cpu_decoding:
test_node_id = test_node_id.cuda()
test_norm = test_norm.cuda()
test_rel = test_rel.cuda()
# validation and testing triplets
valid_data = torch.LongTensor(valid_data)
test_data = torch.LongTensor(test_data)
if use_cuda and not args.no_cuda and not cpu_eval:
valid_data = valid_data.cuda()
test_data = test_data.cuda()
test_graph.ndata.update({'id': test_node_id, 'norm': test_norm})
test_graph.edata['type'] = test_rel
if use_cuda and not args.no_cuda:
# model = nn.DataParallel(model, device_ids=[0,1])
model = model.cuda()
model_state_file = get_model_name(args)
# writer = SummaryWriter("runs/" + model_state_file.replace(".pth",".log"))
# check if only evaluation needs to be performed
if args.eval_only:
if args.load_model:
model_state_file = args.load_model
else:
print("Please provide model path for evaluation (--load_model)")
sys.exit(0)
checkpoint = torch.load(model_state_file)
if use_cuda and not args.no_cuda and cpu_eval:
model.cpu()
test_graph.ndata['id'] = test_graph.ndata['id'].cpu()
test_graph.ndata['norm'] = test_graph.ndata['norm'].cpu()
test_graph.edata['type'] = test_graph.edata['type'].cpu()
model.decoder.no_cuda = True
model.eval()
model.load_state_dict(checkpoint['state_dict'])
print(model)
print("================DEV=================")
mrr = evaluation_utils.ranking_and_hits(
test_graph,
model,
valid_data,
all_e1_to_multi_e2,
train_network,
fusion="graph-only",
sim_relations=args.sim_relations,
write_results=args.write_results,
debug=args.debug)
print("================TEST================")
mrr = evaluation_utils.ranking_and_hits(
test_graph,
model,
test_data,
all_e1_to_multi_e2,
train_network,
fusion="graph-only",
sim_relations=args.sim_relations,
debug=args.debug)
sys.exit(0)
if os.path.isfile(model_state_file):
print(model_state_file)
overwrite = input('Model already exists. Overwrite? Y = yes, N = no\n')
if overwrite.lower() == 'n':
print("Quitting")
sys.exit(0)
elif overwrite.lower() != 'y':
raise ValueError("Invalid Option")
# build adj list and calculate degrees for sampling
adj_list, degrees, sparse_adj_matrix, rel = utils.get_adj_and_degrees(
num_nodes, num_rels, train_data)
# remove sim edges from sampling_edge_ids (we sample from the original graph and then densify it)
if args.sim_relations:
sim_edge_ids = np.where(graph_train_data[:, 1] == num_rels - 1)[0]
sampling_edge_ids = np.delete(np.arange(len(graph_train_data)),
sim_edge_ids)
else:
sampling_edge_ids = None
# optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
forward_time = []
backward_time = []
# training loop
print("Starting training...")
epoch = 0
best_mrr = 0
while True:
model.train()
epoch += 1
cur_train_data = graph_train_data[:]
# build dgl graph
g, node_id, edge_type, node_norm, data, labels = \
utils.generate_sampled_graph_and_labels(
cur_train_data, args.graph_batch_size,
num_rels, adj_list, degrees, args.negative_sample, args.sim_sim, args.sim_relations,
sim_train_e1_to_multi_e2, sampling_edge_ids)
node_id_copy = np.copy(node_id)
node_dict = {v: k for k, v in dict(enumerate(node_id_copy)).items()}
# set node/edge feature
node_id = torch.from_numpy(node_id).view(-1, 1)
edge_type = torch.from_numpy(edge_type)
node_norm = torch.from_numpy(node_norm).view(-1, 1)
if use_cuda and not args.no_cuda:
node_id = node_id.cuda()
edge_type, node_norm = edge_type.cuda(), node_norm.cuda()
g.ndata.update({'id': node_id, 'norm': node_norm})
g.edata['type'] = edge_type
batch_size = args.decoder_batch_size
e1_keys = list(train_e1_to_multi_e2.keys())
sub_e1_keys = {}
# Add inverse edges to training samples
src, dst = np.concatenate((cur_train_data[:, 0], cur_train_data[:, 2])), \
np.concatenate((cur_train_data[:, 2], cur_train_data[:, 0]))
rel = cur_train_data[:, 1]
rel = np.concatenate((rel, rel + num_rels))
cur_train_data = np.stack((src, rel, dst)).transpose()
# The loop below constructs a dict for the decoding step
# with the key (src, rel) and the value as the list of nodes present in the original graph
# where the source and target nodes both belong to the list of sampled nodes in subgraph
for e in cur_train_data:
rel = e[1]
# Don't use sim relations for decoding
if args.sim_relations:
if rel == num_rels - 1 or rel == (num_rels * 2) - 1:
continue
elif rel >= num_rels:
rel -= 1
if e[0] in node_id_copy and e[2] in node_id_copy:
subgraph_src_idx = node_dict[e[0]]
subgraph_tgt_idx = node_dict[e[2]]
if (subgraph_src_idx, rel) not in sub_e1_keys:
sub_e1_keys[(subgraph_src_idx, rel)] = [subgraph_tgt_idx]
else:
sub_e1_keys[(subgraph_src_idx,
rel)].append(subgraph_tgt_idx)
key_list = list(sub_e1_keys.keys())
random.shuffle(key_list)
cum_loss = 0.0
for i in range(0, len(key_list), batch_size):
optimizer.zero_grad()
# compute graph embeddings
graph_embeddings = model.get_graph_embeddings(g, epoch)
#model.decoder.module.cur_embedding = graph_embeddings
model.decoder.cur_embedding = graph_embeddings
batch = key_list[i:i + batch_size]
# Don't train with batches of size 1 and always set batch_size > 1 since batch norm
# fails with batch_size=1
if len(batch) == 1:
continue
e1 = torch.LongTensor([elem[0] for elem in batch])
rel = torch.LongTensor([elem[1] for elem in batch])
# e2 -> list of target nodes in subgraph
e2 = [sub_e1_keys[(k[0], k[1])] for k in batch]
batch_len = len(batch)
if use_cuda and not args.no_cuda and not cpu_decoding:
target = torch.cuda.FloatTensor(batch_len,
node_id_copy.shape[0]).fill_(0)
e1 = e1.cuda()
rel = rel.cuda()
else:
target = torch.zeros((batch_len, node_id_copy.shape[0]))
# construct target tensor
for j, inst in enumerate(e2):
target[j, inst] = 1.0
# perform label smoothing
target = ((1.0 - args.label_smoothing_epsilon) *
target) + (1.0 / target.size(1))
if cpu_decoding:
graph_embeddings = graph_embeddings.cpu()
model.decoder.cpu()
model.decoder.no_cuda = True
t0 = time.time()
loss = model.get_score(e1, rel, target, graph_embeddings)
loss = torch.mean(loss)
cum_loss += loss.cpu().item()
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)
del graph_embeddings, target, batch, loss, e1, rel, e2
# the below make training very slow
# gc.collect()
# torch.cuda.empty_cache()
print(
"Epoch {:04d} | Loss {:.4f} | Best MRR {:.4f} | Forward {:.4f}s | Backward {:.4f}s"
.format(epoch, cum_loss, best_mrr, forward_time[-1],
backward_time[-1]))
# writer.add_scalar('data/loss', cum_loss , epoch)
# Save model every 100 epochs
# if epoch + 1 % 100==0:
# print("saving current model..")
# torch.save({'state_dict': model.state_dict(), 'epoch': epoch},
# model_state_file)
# validation
if epoch % args.evaluate_every == 0:
# perform validation on CPU when full graph is too large
if use_cuda and not args.no_cuda and cpu_eval:
model.cpu()
test_graph.ndata['id'] = test_graph.ndata['id'].cpu()
test_graph.ndata['norm'] = test_graph.ndata['norm'].cpu()
test_graph.edata['type'] = test_graph.edata['type'].cpu()
model.decoder.no_cuda = True
model.eval()
print("start eval")
print("===========DEV============")
mrr = evaluation_utils.ranking_and_hits(
test_graph,
model,
valid_data,
all_e1_to_multi_e2,
train_network,
fusion="graph-only",
sim_relations=args.sim_relations,
debug=args.debug,
epoch=epoch)
# writer.add_scalar('data/mrr', mrr, epoch)
# save best model
# torch.save({'state_dict': model.state_dict(), 'epoch': epoch},
# model_state_file)
if mrr < best_mrr:
if epoch >= args.n_epochs:
break
else:
best_mrr = mrr
print("[saving best model so far]")
torch.save({
'state_dict': model.state_dict(),
'epoch': epoch
}, model_state_file)
metrics = {"best_mrr": best_mrr, "cum_loss": cum_loss}
with open(os.path.join(args.output_dir, 'metrics.json'), 'w') as f:
f.write(json.dumps(metrics))
# transfer graph back to gpu device
if use_cuda and not args.no_cuda and cpu_eval:
model.cuda()
test_graph.ndata['id'] = test_graph.ndata['id'].cuda()
test_graph.ndata['norm'] = test_graph.ndata['norm'].cuda()
test_graph.edata['type'] = test_graph.edata['type'].cuda()
model.decoder.no_cuda = False
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)))
# writer.export_scalars_to_json("./all_scalars.json")
# writer.close()
print("\nStart testing")
# use best model checkpoint
checkpoint = torch.load(model_state_file)
model.eval()
model.load_state_dict(checkpoint['state_dict'])
print("Using best epoch: {}".format(checkpoint['epoch']))
evaluation_utils.ranking_and_hits(test_graph,
model,
test_data,
all_e1_to_multi_e2,
train_network,
fusion="graph-only",
sim_relations=args.sim_relations)
def main(params):
# Load data
dataset = Dataset(params['data_path'], params['test']['test_target'])
num_nodes = dataset.num_nodes
train_data = dataset.train
valid_data = dataset.valid
test_data = dataset.test
num_relations = dataset.num_relations
# validation and testing triplets
valid_data = torch.tensor(valid_data, dtype=torch.long)
test_data = torch.tensor(test_data, dtype=torch.long)
train_data_tensor = torch.tensor(train_data, dtype=torch.long)
# build test graph
test_graph, test_rel, test_norm = utils.build_test_graph(
num_nodes, num_relations, train_data)
test_node_id = torch.arange(0, num_nodes, dtype=torch.long).view(-1, 1)
test_rel = torch.from_numpy(test_rel)
test_norm = utils.node_norm_to_edge_norm(
test_graph, torch.from_numpy(test_norm).view(-1, 1))
# build adj list and calculate degrees for sampling
adj_list, degrees = utils.get_adj_and_degrees(num_nodes, train_data)
# GPU settings
use_cuda = params['use_cuda'] and torch.cuda.is_available()
cpu_device = torch.device('cpu')
if not use_cuda:
device1 = cpu_device
device2 = cpu_device
else:
device1 = torch.device(params['graph_encoder']['device'])
device2 = torch.device(params['relation_decoder']['device'])
logger.info('device 1: {}, device 2: {}'.format(
str(device1), str(device2)))
# Load modules
# Load pre-trained embedding
embed = None
if params['load_embed']['do']:
pretrain_embed = np.load(params['load_embed']['embed_path'])
embed = torch.from_numpy(pretrain_embed['embed'])
logger.info('Loaded pretrained entity embedding from {}'.format(
params['load_embed']['embed_path']))
if 'rel' in pretrain_embed:
rel = pretrain_embed['rel']
logger.info('Loade pretrained relation embedding from {}'.format(
params['load_embed']['embed_path']))
# Load graph encoder and relation decoder
graph_encoder = load_graph_encoder(params, dataset, embed,
test_graph, test_rel, test_norm)
graph_encoder.set_device(device1)
relation_decoder = load_relation_decoder(params, dataset)
relation_decoder.set_device(device2)
# Model path for saving/loading
model_state_file = params['model_path']
learning_rate = params['train']['lr']
weight_decay = params['train']['weight_decay']
optimizer = Adam(list(graph_encoder.parameters()) + list(relation_decoder.parameters()),
lr=learning_rate, weight_decay=weight_decay)
# Training
if params['train']['do']:
logger.info('Start training...')
# Use tensorboard to record scalars
writer = SummaryWriter(params['train']['log_file'])
epoch = 0
n_epochs = params['train']['n_epochs']
best_mrr = 0
while epoch < n_epochs:
epoch += 1
# Prepare model
for model in [graph_encoder, relation_decoder]:
model.train()
optimizer.zero_grad()
# perform edge neighborhood sampling to generate training graph and data
batch_size = params['train']['train_batch_size']
split_size = params['train']['graph_split_size']
negative_sample = params['train']['negative_sample']
edge_sampler = params['train']['edge_sampler']
g, node_id, rel, node_norm, data, labels = \
utils.generate_sampled_graph_and_labels(
train_data, batch_size, split_size,
num_relations, adj_list, degrees,
negative_sample, edge_sampler)
logger.info('Done edge sampling')
# set node / edge feature
node_id = torch.from_numpy(node_id).view(-1, 1).long()
rel = torch.from_numpy(rel)
edge_norm = utils.node_norm_to_edge_norm(
g, torch.from_numpy(node_norm).view(-1, 1))
data, labels = torch.from_numpy(data), torch.from_numpy(labels)
# Forward pass
t0 = time.time()
emb_entity = graph_encoder(g, node_id, rel, edge_norm)
score = relation_decoder(emb_entity, data)
loss = relation_decoder.get_loss(score, labels)
# Add regularization if necessary
# regularization = params['train']['regularization']
# loss += regularization * \
# (torch.pow(emb_entity, 2).mean() +
# relation_decoder.reglurization())
t1 = time.time()
# Record loss
writer.add_scalar('loss', loss.item(), epoch)
# clip gradients
loss.backward()
for model in [graph_encoder, relation_decoder]:
torch.nn.utils.clip_grad_norm_(
model.parameters(), params['train']['grad_norm'])
# Optimize
optimizer.step()
t2 = time.time()
forward_time = t1 - t0
backward_time = t2 - t1
logger.info('Epoch {:04d} | Loss {:.4f} | Best MRR {:.4f} | '
'Forward {:.4f}s | Backward {:.4f}s'.
format(epoch, loss.item(), best_mrr, forward_time, backward_time))
# validation
if epoch % params['train']['eval_every'] == 0:
# perform validation on CPU because full graph is too large
if use_cuda:
graph_encoder.set_device(cpu_device)
relation_decoder.set_device(cpu_device)
logger.info('start eval')
with torch.no_grad():
emb_entity = graph_encoder(
test_graph, test_node_id, test_rel, test_norm)
logger.info(
'graph encoded embedding of each node calculated')
res = utils.eval_filtered(emb_entity, relation_decoder,
train_data_tensor, valid_data, test_data,
hits=[1, 3, 10], eval_type='test')
mrr = res['tail']['mrr']
# Record evaluation results
for rank_name, eval_result in res.items():
for k, v in eval_result.items():
writer.add_scalar(
'{}_{}'.format(rank_name, k), v, epoch)
# save best model
if mrr < best_mrr:
if epoch >= params['train']['n_epochs']:
break
else:
best_mrr = mrr
utils.save_model(model_state_file, graph_encoder,
relation_decoder, epoch)
# Recover device
graph_encoder.set_device(device1)
relation_decoder.set_device(device2)
logger.info('training done')
if params['test']['do']:
logger.info('start testing:')
# use best model checkpoint
epoch = utils.load_model(
model_state_file, graph_encoder, relation_decoder)
if epoch:
logger.info('Restore model from: {}, '
'using best epoch: {}'.format(model_state_file,
epoch))
# perform validation on CPU because full graph is too large
if use_cuda:
graph_encoder.set_device(cpu_device)
relation_decoder.set_device(cpu_device)
with torch.no_grad():
emb_entity = graph_encoder(
test_graph, test_node_id, test_rel, test_norm)
utils.eval_filtered(emb_entity, relation_decoder,
train_data_tensor, valid_data,
test_data, hits=[1, 3, 10],
eval_type='test', entity_filters=dataset.entity_filters)
logger.info('testing done')
if params['export_embed']['do']:
logger.info('exporting embedding from graph model...')
# use best model checkpoint
epoch = utils.load_model(
model_state_file, graph_encoder, relation_decoder)
if epoch:
logger.info('Restore model from: {}, '
'using best epoch: {}'.format(model_state_file,
epoch))
# Save embedding weights
save_file = params['export_embed']['embed_path']
emb_weights = graph_encoder.emb_node.weight.detach().cpu().numpy()
np.savez_compressed(save_file, embed=emb_weights)
logger.info('Saved embedding weights to ' + save_file)
def main(args):
fold = 10
topk = 100
if not os.path.exists('result/'):
os.mkdir('result/')
if not os.path.exists('rules/'):
os.mkdir('rules/')
path = 'dataset/' + args.dataset + '/' + str(fold) + '_fold' + '/'
complete_data, num_nodes, num_rels, node_features = load_whole_data(path, args.ftype, args.n_hidden)
train_data = complete_data
test_data = complete_data
in_feat = node_features.shape[1]
node_features = torch.from_numpy(np.array(node_features)).float()
# check cuda
use_cuda = args.gpu >= 0 and torch.cuda.is_available()
if use_cuda:
torch.cuda.set_device(args.gpu)
if use_cuda:
node_features = node_features.cuda()
model_state_file = args.dataset + '_' + args.ftype + '_all_data_model_state.pth'
model = BinaryPredict(in_feat,
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,
node_features=node_features)
if use_cuda:
model.cuda()
# optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
train_idx = np.arange(len(train_data))
valid_idx = sorted(random.sample(list(train_idx), len(train_idx) // 5))
train_idx = sorted(list(set(train_idx) - set(valid_idx)))
valid_data = train_data[valid_idx]
train_data = train_data[train_idx]
valid_data = torch.LongTensor(valid_data)
test_data = torch.LongTensor(test_data)
# build test graph
test_graph, test_rel, test_norm, edge_norm = utils.build_test_graph(
num_nodes, num_rels, complete_data)
test_node_id = torch.arange(0, num_nodes, dtype=torch.long)
test_rel = torch.from_numpy(test_rel)
edge_norm = torch.from_numpy(edge_norm).unsqueeze(1)
test_rel = test_rel.long()
test_norm = torch.from_numpy(test_norm)
if use_cuda:
test_node_id = test_node_id.cuda()
test_rel, test_norm = test_rel.cuda(), test_norm.cuda()
edge_norm = edge_norm.cuda()
test_graph.ndata.update({'id': test_node_id, 'norm': test_norm}) # 'id': test_node_id,
test_graph.edata.update({'type': test_rel, 'norm': edge_norm})
# build adj list and calculate degrees for sampling
adj_list, degrees = utils.get_adj_and_degrees(num_nodes, train_data)
# training loop
print("start training...")
epoch = 0
best_f1 = 0
while True:
model.train()
epoch += 1
g, node_id, edge_type, node_norm, data, labels, edge_norm = \
utils.generate_sampled_graph_and_labels(
train_data, args.graph_batch_size, args.graph_split_size,
num_rels, adj_list, degrees, args.negative_sample)
node_id = torch.from_numpy(node_id).long()
edge_type = torch.from_numpy(edge_type)
edge_type = edge_type.long()
node_norm = torch.from_numpy(node_norm)
data, labels = torch.from_numpy(data), torch.from_numpy(labels)
deg = g.in_degrees(range(g.number_of_nodes())).float()
edge_norm = torch.from_numpy(edge_norm).unsqueeze(1)
if use_cuda:
node_id, deg = node_id.cuda(), deg.cuda()
edge_type, node_norm, edge_norm = edge_type.cuda(), node_norm.cuda(), edge_norm.cuda()
data, labels = data.cuda(), labels.cuda()
g.ndata.update({'id': node_id, 'norm': node_norm})
g.edata['type'] = edge_type
g.edata['norm'] = edge_norm
loss = model.get_loss(g, data, labels)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_norm) # clip gradients
optimizer.step()
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()
scores_val, labels_val = utils.compute_score(test_graph, model, valid_data) # predicted probability
scores_val = scores_val.detach().numpy()
best_thred, f1_val = utils._select_threshold(labels_val, scores_val)
if f1_val < best_f1:
if epoch >= args.n_epochs:
break
else:
best_f1 = f1_val
torch.save({'state_dict': model.state_dict(), 'epoch': epoch, 'best_threshold': best_thred},
model_state_file)
if use_cuda:
model.cuda()
print("training done")
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']))
scores_test, labels_test = utils.compute_score(test_graph, model, test_data)
scores_test = np.array(scores_test)
labels_test = np.array(labels_test)
pred_rules = utils.find_rules_bt_all(path, test_data, labels_test, scores_test, topk)
file_name = args.ftype + '_' + args.dataset + '_all_data.txt'
f = open('rules/pred_' + file_name, 'w', encoding='utf-8')
f.write(pred_rules)
f.close()
def main(args):
# load graph data
data = knwlgrh.load_link(args.dataset)
num_nodes = data.num_nodes
train_data = data.train # np.array (n, 3)
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,
rel_type=args.rel_type,
num_bases=args.n_bases,
num_hidden_layers=args.n_layers,
dropout=args.dropout,
use_cuda=use_cuda,
reg_param=args.regularization)
for foo in model.parameters():
print(foo.size())
# build test graph
test_graph, test_rel, test_norm = utils.build_test_graph(
num_nodes, num_rels, train_data)
test_node_id = torch.arange(0, num_nodes, dtype=torch.long).view(-1, 1)
test_rel = torch.from_numpy(test_rel)
test_norm = torch.from_numpy(test_norm).view(-1, 1)
test_graph.ndata.update({'id': test_node_id, 'norm': test_norm})
# edges have unique IDs according to order of population
# labels each edge with the relation type
test_graph.edata['type'] = test_rel
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)
# training loop
print("start training...")
epoch = 0
while True:
model.train()
epoch += 1
# perform edge neighborhood sampling to generate training graph and data
# return g, uniq_v, rel, norm, samples, labels
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.loss)
# only half of the sampled data is used to create g
# but "data" is all of the sampled data
print("Done edge sampling")
# set node/edge feature
node_id = torch.from_numpy(node_id).view(-1, 1)
edge_type = torch.from_numpy(edge_type)
node_norm = 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, node_norm = edge_type.cuda(), node_norm.cuda()
data, labels = data.cuda(), labels.cuda()
# assigns node IDs
g.ndata.update({'id': node_id, 'norm': node_norm})
# assigns node types
g.edata['type'] = edge_type
# perform propagation using half the sampled data only (g)
# but evaluate the learned embedding using all of the sampled data (data)
if args.loss == "bce":
loss = model.bce_loss(g, data, labels)
elif args.loss == "cos":
loss = model.cos_loss(g, data, labels)
else:
raise NameError("loss function not specified")
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(),
args.grad_norm) # clip gradients
optimizer.step()
eprint("Epoch {:04d} | Loss {:.4f}".format(epoch, loss.item()))
optimizer.zero_grad()
if epoch % args.save_every == 0:
node_vec = model(test_graph)
torch.save(
node_vec, args.dataset + "_" + args.rel_type + "rel_" +
args.loss + "loss_" + str(args.n_hidden) + "dim_" +
str(epoch) + "epochs_nodes" + ".pt")
torch.save(
list(model.parameters())[0].data,
args.dataset + "_" + args.rel_type + "rel_" + args.loss +
"loss_" + str(args.n_hidden) + "dim_" + str(epoch) +
"epochs_rels" + ".pt")
if epoch == args.n_epochs:
break
print("Training finished")
def main(args):
if not os.path.exists('result/'):
os.mkdir('result/')
if not os.path.exists('rules/'):
os.mkdir('rules/')
results = []
fold = 10
results_out = str(args) + '\n'
true_rules_out = ''
pred_rules_out = ''
for i in range(fold):
path = 'dataset/' + args.dataset + '/' + str(fold) + '_fold/'
train_data, test_data, num_nodes, num_rels, node_features = train_test_triples(
path, i, args.ftype, args.n_hidden)
complete_data = np.concatenate((train_data, test_data))
print("Edges: ", len(complete_data))
if test_data.shape[0] == 0:
continue
# create model
in_feat = node_features.shape[1]
node_features = torch.from_numpy(np.array(node_features)).float()
# check cuda
use_cuda = args.gpu >= 0 and torch.cuda.is_available()
if use_cuda:
torch.cuda.set_device(args.gpu)
if use_cuda:
node_features = node_features.cuda()
model_state_file = args.dataset + '_' + args.ftype + '_model_state.pth'
model = BinaryPredict(in_feat,
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,
node_features=node_features,
relation_features=None)
if use_cuda:
model.cuda()
# optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
# validation and testing triples
train_idx = np.arange(len(train_data))
valid_idx = sorted(random.sample(list(train_idx), len(train_idx) // 5))
train_idx = sorted(list(set(train_idx) - set(valid_idx)))
valid_data = train_data[valid_idx]
train_data = train_data[train_idx]
valid_data = torch.LongTensor(valid_data)
test_data = torch.LongTensor(test_data)
# build test graph
test_graph, test_rel, test_norm, edge_norm = utils.build_test_graph(
num_nodes, num_rels, complete_data)
test_node_id = torch.arange(0, num_nodes, dtype=torch.long)
test_rel = torch.from_numpy(test_rel)
edge_norm = torch.from_numpy(edge_norm).unsqueeze(1)
test_rel = test_rel.long()
test_norm = torch.from_numpy(test_norm)
if use_cuda:
test_node_id = test_node_id.cuda()
test_rel, test_norm = test_rel.cuda(), test_norm.cuda()
edge_norm = edge_norm.cuda()
test_graph.ndata.update({
'id': test_node_id,
'norm': test_norm
}) #'id': test_node_id,
test_graph.edata.update({'type': test_rel, 'norm': edge_norm})
# build adj list and calculate degrees for sampling
adj_list, degrees = utils.get_adj_and_degrees(num_nodes, train_data)
# training loop
print("start training...")
epoch = 0
best_f1 = 0
while True:
model.train()
epoch += 1
g, node_id, edge_type, node_norm, data, labels, edge_norm = \
utils.generate_sampled_graph_and_labels(
train_data, args.graph_batch_size, args.graph_split_size,
num_rels, adj_list, degrees, args.negative_sample)
node_id = torch.from_numpy(node_id).long()
edge_type = torch.from_numpy(edge_type)
edge_type = edge_type.long()
node_norm = torch.from_numpy(node_norm)
data, labels = torch.from_numpy(data), torch.from_numpy(labels)
deg = g.in_degrees(range(g.number_of_nodes())).float()
edge_norm = torch.from_numpy(edge_norm).unsqueeze(1)
if use_cuda:
node_id, deg = node_id.cuda(), deg.cuda()
edge_type, node_norm, edge_norm = edge_type.cuda(
), node_norm.cuda(), edge_norm.cuda()
data, labels = data.cuda(), labels.cuda()
g.ndata.update({'id': node_id, 'norm': node_norm})
g.edata['type'] = edge_type
g.edata['norm'] = edge_norm
loss = model.get_loss(g, data, labels)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(),
args.grad_norm) # clip gradients
optimizer.step()
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()
scores_val, labels_val = utils.compute_score(
test_graph, model, valid_data) # predicted probability
scores_val = scores_val.detach().numpy()
best_thred, f1_val = utils._select_threshold(
labels_val, scores_val)
if f1_val < best_f1:
if epoch >= args.n_epochs:
break
else:
best_f1 = f1_val
torch.save(
{
'state_dict': model.state_dict(),
'epoch': epoch,
'best_threshold': best_thred
}, model_state_file)
if use_cuda:
model.cuda()
print("training done")
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']))
best_thred = checkpoint['best_threshold']
scores_test, labels_test = utils.compute_score(test_graph, model,
test_data)
scores_test = scores_test.detach().numpy()
labels_test_pred = (scores_test >
best_thred) * np.ones_like(scores_test)
test_rules, tmp1 = utils.find_rules_bt(path, i, test_data, labels_test)
pred_rules, tmp2 = utils.find_rules_bt(path, i, test_data,
labels_test_pred)
true_rules_out += 'fold ' + str(i) + '\n'
true_rules_out += tmp1 + '\n'
pred_rules_out += 'fold ' + str(i) + '\n'
pred_rules_out += tmp2 + '\n'
precision_test, recall_test, f1_test = utils.metrics(
test_rules, pred_rules)
results_out += "Test Precision: {:.4f} | Test Recall: {:.4f} | Test F1: {:.4f} \n".format(
precision_test, recall_test, f1_test)
print(
"Test Precision: {:.4f} | Test Recall: {:.4f} | Test F1: {:.4f} \n"
.format(precision_test, recall_test, f1_test))
results.append([precision_test, recall_test, f1_test])
mean_p, mean_r, mean_f1 = np.mean(np.array(results), axis=0)
results_out += "Mean values over " + str(
fold
) + " fold: Precision: {:.4f} | Recall: {:.4f} | F1: {:.4f}".format(
mean_p, mean_r, mean_f1)
print("Mean values over " + str(fold) +
" fold: Precision: {:.4f} | Recall: {:.4f} | F1: {:.4f}".format(
mean_p, mean_r, mean_f1))
file_name = args.ftype + '_' + args.dataset + '.txt'
f = open('result/' + file_name, 'w', encoding='utf-8')
f.write(results_out)
f.close()
f = open('rules/true_' + file_name, 'w', encoding='utf-8')
f.write(true_rules_out)
f.close()
f = open('rules/pred_' + file_name, 'w', encoding='utf-8')
f.write(pred_rules_out)
f.close()
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
all_data = np.concatenate((train_data, valid_data, test_data), axis=0)
if (use_shuriken):
monitor = ShurikenMonitor()
use_cuda = torch.cuda.is_available()
# create model
model = LinkPredict(num_nodes,
args.n_hidden,
num_rels,
args.model,
num_bases=args.n_bases,
num_hidden_layers=args.n_layers,
dropout=args.dropout,
use_cuda=use_cuda,
reg_param=args.regularization,
skip_connection=args.skip_connection,
rel_activation=args.rel_activation,
rel_dropout=args.rel_dropout)
# check if there is a model with the same hyperparameters saved
new_model, res = load_model(args, model)
epoch = 0
if (res != 0):
model = new_model
epoch = res
best_model, best_mrr = load_best_model(args, model)
# validation and testing triplets
valid_data = torch.LongTensor(valid_data)
test_data = torch.LongTensor(test_data)
# all_data = torch.LongTensor(all_data.astype(set))
# build test graph
test_graph, test_rel, test_norm, incidence_in_test, incidence_out_test = 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 = torch.from_numpy(test_norm).view(-1, 1)
test_graph.ndata.update({'id': test_node_id, 'norm': test_norm})
test_graph.edata['type'] = test_rel
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_' + str(args.dropout) + '_' + str(args.lr) + '_' + str(args.regularization) + '.pth'
forward_time = []
backward_time = []
print("start training...")
weight = torch.randn(num_rels * 2, args.n_hidden).cuda()
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, incidence_in, incidence_out = \
utils.generate_sampled_graph_and_labels(
train_data, args.graph_batch_size, args.graph_split_size,
num_rels, adj_list, degrees, args.negative_sample)
# set node/edge feature
node_id = torch.from_numpy(node_id).view(-1, 1)
edge_type = torch.from_numpy(edge_type)
node_norm = 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, node_norm = edge_type.cuda(), node_norm.cuda()
data, labels = data.cuda(), labels.cuda()
g.ndata.update({'id': node_id, 'norm': node_norm})
g.edata['type'] = edge_type
t0 = time.time()
loss, weight = model.get_loss(g, data, labels, weight, incidence_in,
incidence_out)
weight = weight.detach()
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:
save_model(args, model, epoch, optimizer)
# perform validation on CPU because full graph is too large
if use_cuda:
model.cpu()
model.eval()
print("start eval")
# mrr_f = utils.evaluate_filtered(test_graph, model, valid_data, all_data, num_nodes, weight, incidence_in_test, incidence_out_test, hits=[1, 3, 10], eval_bz=args.eval_batch_size)
mrr = utils.evaluate(test_graph,
model,
valid_data,
num_nodes,
weight,
incidence_in_test,
incidence_out_test,
hits=[1, 3, 10],
eval_bz=args.eval_batch_size)
if (use_shuriken):
monitor.send_info(epoch, {"mrr": mrr})
# save best model
if mrr < best_mrr:
if epoch >= args.n_epochs:
break
else:
best_mrr = mrr
best_model = model
try:
os.makedirs(args.model)
except FileExistsError:
pass
save_best_model(args, model, best_mrr)
if use_cuda:
model.cuda()
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(args.model + "/" + model_state_file)
if use_cuda:
best_model.cpu() # test on CPU
best_model.eval()
test_mrr = utils.evaluate(test_graph,
best_model,
test_data,
num_nodes,
weight,
incidence_in_test,
incidence_out_test,
hits=[1, 3, 10],
eval_bz=args.eval_batch_size)
if (use_shuriken):
monitor.send_info(epoch, {"test mrr": test_mrr})
epoch = 0
best_mrr = 0
while epoch < 10000:
model.train()
epoch += 1
# perform edge neighborhood sampling to generate training graph and data
index = random.randint(250, 450)
print(index)
print(len(triplets[index]))
g, node_id, edge_type, node_norm, data, labels = \
utils.generate_sampled_graph_and_labels(
triplets[index], 1000, 0.5,
num_rels, adj_lists[index], degrees_lists[index], 100)
print("Done edge sampling")
# set node/edge feature
node_id = torch.from_numpy(node_id).view(-1, 1)
edge_type = torch.from_numpy(edge_type)
node_norm = 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)
g.ndata.update({'id': node_id, 'norm': node_norm})
g.edata['type'] = edge_type
t0 = time.time()
loss = model.get_loss(g, data, labels)
t1 = time.time()
def main(args):
# initialize visualization environment as global variable
# vis = VisManger('main')
vis = ''
# initialize the data importer and load data
kg = RGCNDataSet(args.directory, args.train,
args.valid, args.test, args.parser)
kg.load_data()
# get dataset data
num_nodes = kg.num_nodes
train_data = kg.train
valid_data = kg.valid
test_data = kg.test
num_rels = kg.num_rels
entity_dict = kg.entity_dict
relation_dict = kg.relation_dict
# check cuda
use_cuda = args.gpu >= 0 and torch.cuda.is_available()
if use_cuda:
torch.cuda.set_device(args.gpu)
# set the use of cuda also for the evaluation
use_cuda_eval = args.gpu_eval >= 0 and torch.cuda.is_available()
# 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 = torch.from_numpy(test_norm).view(-1, 1)
if use_cuda_eval:
test_deg = test_deg.cuda()
test_node_id = test_node_id.cuda()
test_rel = test_rel.cuda()
test_norm = test_norm.cuda()
test_graph.ndata.update({'id': test_node_id, 'norm': test_norm})
test_graph.edata['type'] = test_rel
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 = args.directory + 'model_datasets/' + 'model_state.pth'
forward_time = []
backward_time = []
# training loop
print("\n *** Start training... ***")
epoch = 0
best_mrr = 0
while True:
model.train()
epoch += 1
# perform edge neighborhood sampling to generate training graph and data
start_sampling_time = time.time()
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)
print("Done edge sampling")
print("Sampling time: %s seconds" %
(time.time() - start_sampling_time))
# set node/edge feature
node_id = torch.from_numpy(node_id).view(-1, 1).long()
edge_type = torch.from_numpy(edge_type)
node_norm = 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, node_norm = edge_type.cuda(), node_norm.cuda()
data, labels = data.cuda(), labels.cuda()
g.ndata.update({'id': node_id, 'norm': node_norm})
g.edata['type'] = edge_type
t0 = time.time()
loss = model.get_loss(g, 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("\nEpoch {:04d} | Loss {:.4f} | Best MRR {:.4f} | Forward {:.4f}s | Backward {:.4f}s".
format(epoch, loss.item(), best_mrr, forward_time[-1], backward_time[-1]))
#vis.plot_loss(np.array(loss.item()), epoch)
optimizer.zero_grad()
# validation
if epoch % args.evaluate_every == 0:
start_valid_time = time.time()
# perform validation on CPU if the full graph is too large, otherwise perform it on the GPU
if use_cuda_eval:
model.cuda()
valid_data = valid_data.cuda()
else:
model.cpu()
model.eval()
print("\n\n*** Perform the evaluation on the validation dataset ***")
print("Epoch: " + str(epoch))
mrr = utils.evaluate(test_graph,
model,
valid_data,
epoch,
entity_dict,
relation_dict,
vis,
args.directory,
args.score,
hits=[1, 3, 10],
eval_bz=args.eval_batch_size)
print("Evaluation time for the validation set: %s seconds" %
(time.time() - start_valid_time))
print("\n")
# save the best model
if mrr < best_mrr:
if epoch >= args.n_epochs:
break
else:
best_mrr = mrr
torch.save({
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'loss': loss
}, model_state_file)
# used for testing reason to interrupt (default value is False)
if args.forced_stop == True:
print("Force stop!")
break
if use_cuda:
model.cuda()
print("*** Training is completed!!! ***")
print("Mean forward time: {:4f}s".format(np.mean(forward_time)))
print("Mean Backward time: {:4f}s".format(np.mean(backward_time)))
print("\n\n*** Perform the evaluation on the test dataset *** ")
# use best model checkpoint
checkpoint = torch.load(model_state_file)
# perform validation on CPU if the full graph is too large, otherwise perform it on the GPU
if use_cuda_eval:
model.cuda()
test_data = test_data.cuda()
else:
model.cpu()
model.eval()
model.load_state_dict(checkpoint['model_state_dict'])
print("Using best epoch: {}".format(checkpoint['epoch']))
start_test_time = time.time()
utils.evaluate(test_graph,
model,
test_data,
checkpoint['epoch'],
entity_dict,
relation_dict,
vis,
args.directory,
'best',
hits=[1, 3, 10],
eval_bz=args.eval_batch_size)
print("Evaluation time for the test set: %s seconds" %
(time.time() - start_test_time))
