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): # 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): 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(args): # load graph data print(time.strftime("%a, %d %b %Y %H:%M:%S +0000: ", time.localtime()) + 'start loading...', flush=True) if args.supervised == 'True': train_pool, train_labels, nlabels, multi = utils.load_label(args.label) train_data, num_nodes, num_rels, train_indices, ntrain, node_attri = utils.load_supervised( args, args.link, args.node, train_pool) elif args.supervised == 'False': train_data, num_nodes, num_rels, node_attri = utils.load_unsupervised( args, args.link, args.node) nlabels = 0 print(time.strftime("%a, %d %b %Y %H:%M:%S +0000: ", time.localtime()) + 'finish loading...', flush=True) # check cuda use_cuda = args.gpu >= 0 and torch.cuda.is_available() if use_cuda: torch.cuda.set_device(args.gpu) print('check 1', flush=True) # create model model = TrainModel(node_attri, num_nodes, args.n_hidden, num_rels, nlabels, num_bases=args.n_bases, num_hidden_layers=args.n_layers, dropout=args.dropout, use_cuda=use_cuda, reg_param=args.regularization) print('check 2', flush=True) if use_cuda: model.cuda() print('check 3', flush=True) # build adj list and calculate degrees for sampling degrees = utils.get_adj_and_degrees(num_nodes, train_data) print('check 4', flush=True) # optimizer optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) # training loop print(time.strftime("%a, %d %b %Y %H:%M:%S +0000: ", time.localtime()) + "start training...", flush=True) for epoch in range(args.n_epochs): model.train() # perform edge neighborhood sampling to generate training graph and data if args.supervised == 'True': g, node_id, edge_type, node_norm, matched_labels, matched_index = \ utils.generate_sampled_graph_and_labels_supervised( train_data, args.graph_batch_size, args.graph_split_size, num_rels, degrees, args.negative_sample, args.edge_sampler, train_indices, train_labels, multi, nlabels, ntrain, if_train=True, label_batch_size=args.label_batch_size) if multi: matched_labels = torch.from_numpy(matched_labels).float() else: matched_labels = torch.from_numpy(matched_labels).long() elif args.supervised == 'False': g, node_id, edge_type, node_norm, data, labels = \ utils.generate_sampled_graph_and_labels_unsupervised( train_data, args.graph_batch_size, args.graph_split_size, num_rels, degrees, args.negative_sample, args.edge_sampler) data, labels = torch.from_numpy(data), torch.from_numpy(labels) # 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)) deg = g.in_degrees(range(g.number_of_nodes())).float().view(-1, 1) if use_cuda: node_id, deg, g = node_id.cuda(), deg.cuda(), g.to('cuda') edge_type, edge_norm = edge_type.cuda(), edge_norm.cuda() if args.supervised == 'True': matched_labels = matched_labels.cuda() elif args.supervised == 'False': data, labels = data.cuda(), labels.cuda() embed, pred = model(g, node_id, edge_type, edge_norm) if args.supervised == 'True': loss = model.get_supervised_loss(pred, matched_labels, matched_index, multi) elif args.supervised == 'False': loss = model.get_unsupervised_loss(g, embed, data, labels) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_norm) # clip gradients optimizer.step() optimizer.zero_grad() print( time.strftime("%a, %d %b %Y %H:%M:%S +0000: ", time.localtime()) + "Epoch {:05d} | Loss {:.4f}".format(epoch, loss.item()), flush=True) print(time.strftime("%a, %d %b %Y %H:%M:%S +0000: ", time.localtime()) + "training done", flush=True) print(time.strftime("%a, %d %b %Y %H:%M:%S +0000: ", time.localtime()) + "start output...", flush=True) model.eval() if args.attributed == 'True': np.random.shuffle(train_data) node_emb, node_over = np.zeros((num_nodes, args.n_hidden)), set() batch_total = math.ceil(len(train_data) / args.graph_batch_size) for batch_num in range(batch_total): # perform edge neighborhood sampling to generate training graph and data g, old_node_id, edge_type, node_norm, data, labels = \ utils.generate_sampled_graph_and_labels_unsupervised( train_data, args.graph_batch_size, args.graph_split_size, num_rels, degrees, args.negative_sample, args.edge_sampler) # set node/edge feature node_id = torch.from_numpy(old_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)) if use_cuda: node_id, g = node_id.cuda(), g.to('cuda') edge_type, edge_norm = edge_type.cuda(), edge_norm.cuda() embed, _ = model(g, node_id, edge_type, edge_norm) node_emb[old_node_id] = embed.detach().cpu().numpy().astype( np.float32) for each in old_node_id: node_over.add(each) print(time.strftime("%a, %d %b %Y %H:%M:%S +0000: ", time.localtime()) + f'finish output batch nubmer {batch_num} -> {batch_total}', flush=True) utils.save(args, node_emb) elif args.attributed == 'False': utils.save( args, model.rgcn.layers[0].embedding.weight.detach().cpu().numpy()) emb, labs = creoSpazio("./data/PubMed/emb.dat") new_emb = TSNEImpl(emb) drawImpl(new_emb, labs, "./ciao1.png") return
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(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): # 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})
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 """ # optimizer triplets = [(edge_src[i], edge_types_[i], edge_dst[i]) for i in range(len(gs))] triplets = [[(triple[0][i], triple[1][i], triple[2][i]) for i in range(len(triple[0]))] for triple in triplets] adj_lists = [] degrees_lists = [] for triple in triplets: adj_list, degrees = utils.get_adj_and_degrees(max_nodes, triple) adj_lists.append(adj_list) degrees_lists.append(degrees) 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
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))