def test_gae(): model = GAE(encoder=lambda x: x) model.reset_parameters() x = torch.Tensor([[1, -1], [1, 2], [2, 1]]) z = model.encode(x) assert z.tolist() == x.tolist() adj = model.decoder.forward_all(z) assert adj.tolist() == torch.sigmoid( torch.Tensor([[+2, -1, +1], [-1, +5, +4], [+1, +4, +5]])).tolist() edge_index = torch.tensor([[0, 1], [1, 2]]) value = model.decode(z, edge_index) assert value.tolist() == torch.sigmoid(torch.Tensor([-1, 4])).tolist() edge_index = torch.tensor([[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]]) data = Data(edge_index=edge_index) data.num_nodes = edge_index.max().item() + 1 data = train_test_split_edges(data, val_ratio=0.2, test_ratio=0.3) z = torch.randn(11, 16) loss = model.recon_loss(z, data.train_pos_edge_index) assert loss.item() > 0 auc, ap = model.test(z, data.val_pos_edge_index, data.val_neg_edge_index) assert auc >= 0 and auc <= 1 and ap >= 0 and ap <= 1
def test_gae(): model = GAE(encoder=lambda x: x) model.reset_parameters() x = torch.Tensor([[1, -1], [1, 2], [2, 1]]) z = model.encode(x) assert z.tolist() == x.tolist() adj = model.decode(z) assert adj.tolist() == torch.sigmoid( torch.Tensor([[+2, -1, +1], [-1, +5, +4], [+1, +4, +5]])).tolist() edge_index = torch.tensor([[0, 1], [1, 2]]) value = model.decode_indices(z, edge_index) assert value.tolist() == torch.sigmoid(torch.Tensor([-1, 4])).tolist() edge_index = torch.tensor([[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]]) data = Data(edge_index=edge_index) data = model.split_edges(data, val_ratio=0.2, test_ratio=0.3) assert data.val_pos_edge_index.size() == (2, 2) assert data.val_neg_edge_index.size() == (2, 2) assert data.test_pos_edge_index.size() == (2, 3) assert data.test_neg_edge_index.size() == (2, 3) assert data.train_pos_edge_index.size() == (2, 5) assert data.train_neg_adj_mask.size() == (11, 11) assert data.train_neg_adj_mask.sum().item() == (11**2 - 11) / 2 - 4 - 6 - 5 z = torch.randn(11, 16) loss = model.recon_loss(z, data.train_pos_edge_index) assert loss.item() > 0 auc, ap = model.test(z, data.val_pos_edge_index, data.val_neg_edge_index) assert auc >= 0 and auc <= 1 and ap >= 0 and ap <= 1
def test_gae(): model = GAE(encoder=lambda x: x) model.reset_parameters() x = torch.Tensor([[1, -1], [1, 2], [2, 1]]) z = model.encode(x) assert z.tolist() == x.tolist() adj = model.decoder.forward_all(z) assert adj.tolist() == torch.sigmoid( torch.Tensor([[+2, -1, +1], [-1, +5, +4], [+1, +4, +5]])).tolist() edge_index = torch.tensor([[0, 1], [1, 2]]) value = model.decode(z, edge_index) assert value.tolist() == torch.sigmoid(torch.Tensor([-1, 4])).tolist() edge_index = torch.tensor([[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]]) data = Data(edge_index=edge_index, num_nodes=11) transform = RandomLinkSplit(split_labels=True, add_negative_train_samples=False) train_data, val_data, test_data = transform(data) z = torch.randn(11, 16) loss = model.recon_loss(z, train_data.pos_edge_label_index) assert loss.item() > 0 auc, ap = model.test(z, val_data.pos_edge_label_index, val_data.neg_edge_label_index) assert auc >= 0 and auc <= 1 and ap >= 0 and ap <= 1
def train(dataset, args, writer = None): task = args.task test_loader = loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True) if task == 'link': model = GAE(models.GNNStack(dataset.num_node_features, args.hidden_dim, int(dataset.num_classes), args)) elif task == 'node': model = models.GNNStack(dataset.num_node_features, args.hidden_dim, int(dataset.num_classes), args) else: raise RuntimeError("Unknown task.") metrics_for_labels = True if args.metrics_for_labels == 'True' else False scheduler, opt = build_optimizer(args, model.parameters()) print("Training \nModel: {}, Data representation: {}. Dataset: {}, Task type: {}". format(args.model_name, args.graph_type, args.dataset, args.task)) metric_text = 'test accuracy' if task == 'node' else 'test precision' for epoch in range(args.epochs): total_loss = 0 model.train() for batch in loader: opt.zero_grad() if task == 'node': pred = model(batch) label = batch.y pred = pred[batch.train_mask] label = label[batch.train_mask] loss = model.loss(pred, label) else: train_pos_edge_index = batch.train_pos_edge_index z = model.encode(batch) loss = model.recon_loss(z, train_pos_edge_index) loss.backward() opt.step() total_loss += loss.item() * batch.num_graphs total_loss /= len(loader.dataset) if writer == None: print(total_loss) else: writer.add_scalar("loss", total_loss, epoch) if epoch % 10 == 0: test_metric, _ = test(loader, model, task = task) if writer == None: print(test_metric, metric_text) else: writer.add_scalar(metric_text, test_metric, epoch) if metrics_for_labels == True and epoch == args.epochs -1: _, labels_metrics = test(loader, model, task = task, metrics_for_labels=metrics_for_labels) print('{} for labels:\n {}'.format(metric_text, labels_metrics))
def run_experiment(args): """ Performing experiment for the given arguments """ dataset, data = load_data(args.dataset) device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # Define Model encoder = create_encoder(args.model, dataset.num_features, args.latent_dim).to(device) decoder = create_decoder(args.decoder).to(device) if args.model == 'GAE': model = GAE(encoder=encoder, decoder=decoder).to(device) else: model = VGAE(encoder=encoder, decoder=decoder).to(device) # Split edges of a torch_geometric.data.Data object into pos negative train/val/test edges # default ratios of positive edges: val_ratio=0.05, test_ratio=0.1 print("Data.edge_index.size", data.edge_index.size(1)) data = model.split_edges(data) node_features, train_pos_edge_index = data.x.to( device), data.train_pos_edge_index.to(device) optimizer = torch.optim.Adam(model.parameters(), lr=0.01) def train_epoch(): """ Performing training over a single epoch and optimize over loss :return: log - loss of training loss """ # Todo: Add logging of results model.train() optimizer.zero_grad() # Compute latent embedding Z latent_embeddings = model.encode(node_features, train_pos_edge_index) # Calculate loss and loss = model.recon_loss(latent_embeddings, train_pos_edge_index) if args.model in ['VGAE']: loss = loss + (1 / data.num_nodes) * model.kl_loss() # Compute gradients loss.backward() # Perform optimization step optimizer.step() # print("Train-Epoch: {} Loss: {}".format(epoch, loss)) # ToDo: Add logging via Tensorboard log = {'loss': loss} return log def test(pos_edge_index, neg_edge_index): model.eval() with torch.no_grad(): # compute latent var z = model.encode(node_features, train_pos_edge_index) # model.test return - AUC, AP return model.test(z, pos_edge_index, neg_edge_index) def test_naive_graph(z, sample_size=1000): if args.sample_dense_evaluation: graph_type = "sampled" z_sample, index_mapping = sample_graph(z, sample_size) t = time.time() adjacency = model.decoder.forward_all( z_sample, sigmoid=(args.decoder == 'dot')) else: graph_type = "full" t = time.time() adjacency = model.decoder.forward_all( z, sigmoid=(args.decoder == 'dot')) print(f"Computing {graph_type} graph took {time.time() - t} seconds.") print( f"Adjacency matrix takes {adjacency.element_size() * adjacency.nelement() / 10 ** 6} MB of memory." ) if args.min_sim_absolute_value is None: args.min_sim_absolute_value, _ = sample_percentile( args.min_sim, adjacency, dist_measure=args.decoder, sample_size=sample_size) if args.sample_dense_evaluation: precision, recall = sampled_dense_precision_recall( data, adjacency, index_mapping, args.min_sim_absolute_value) else: precision, recall = dense_precision_recall( data, adjacency, args.min_sim_absolute_value) print("Predicted {} adjacency matrix has precision {} and recall {}!". format(graph_type, precision, recall)) return precision, recall def sample_graph(z, sample_size): N, D = z.shape sample_size = min(sample_size, N) sample_ix = np.random.choice(np.arange(N), size=sample_size, replace=False) # Returns the sampled embeddings, and a mapping from their indices to the originals return z[sample_ix], {i: sample_ix[i] for i in np.arange(sample_size)} def test_compare_lsh_naive_graphs(z, assure_correctness=True): """ :param z: :param assure_correctness: :return: """ # Naive Adjacency-Matrix (Non-LSH-Version) t = time.time() # Don't use sigmoid in order to directly compare thresholds with LSH naive_adjacency = model.decoder.forward_all( z, sigmoid=(args.decoder == 'dot')) naive_time = time.time() - t naive_size = naive_adjacency.element_size() * naive_adjacency.nelement( ) / 10**6 if args.min_sim_absolute_value is None: args.min_sim_absolute_value, _ = sample_percentile( args.min_sim, z, dist_measure=args.decoder) print( "______________________________Naive Graph Computation KPI____________________________________________" ) print(f"Computing naive graph took {naive_time} seconds.") print(f"Naive adjacency matrix takes {naive_size} MB of memory.") # LSH-Adjacency-Matrix: t = time.time() lsh_adjacency = LSHDecoder(bands=args.lsh_bands, rows=args.lsh_rows, verbose=True, assure_correctness=assure_correctness, sim_thresh=args.min_sim_absolute_value)(z) lsh_time = time.time() - t lsh_size = lsh_adjacency.element_size() * lsh_adjacency._nnz() / 10**6 print( "__________________________________LSH Graph Computation KPI__________________________________________" ) print(f"Computing LSH graph took {lsh_time} seconds.") print(f"Sparse adjacency matrix takes {lsh_size} MB of memory.") print( "________________________________________Precision-Recall_____________________________________________" ) # 1) Evaluation: Both Adjacency matrices against ground truth graph naive_precision, naive_recall = dense_precision_recall( data, naive_adjacency, args.min_sim_absolute_value) lsh_precision, lsh_recall = sparse_precision_recall( data, lsh_adjacency) print( f"Naive-Precision {naive_precision}; Naive-Recall {naive_recall}") print(f"LSH-Precision {lsh_precision}; LSH-Recall {lsh_recall}") print( "_____________________________Comparison Sparse vs Dense______________________________________________" ) # 2) Evation: Compare both adjacency matrices against each other compare_precision, compare_recall = sparse_v_dense_precision_recall( naive_adjacency, lsh_adjacency, args.min_sim_absolute_value) print( f"LSH sparse matrix has {compare_precision} precision and {compare_recall} recall w.r.t. the naively generated dense matrix!" ) return naive_precision, naive_recall, naive_time, naive_size, lsh_precision, lsh_recall, lsh_time, lsh_size, compare_precision, compare_recall # Training routine early_stopping = EarlyStopping(args.use_early_stopping, patience=args.early_stopping_patience, verbose=True) logs = [] if args.load_model and os.path.isfile("checkpoint.pt"): print("Loading model from savefile...") model.load_state_dict(torch.load("checkpoint.pt")) if not (args.load_model and args.early_stopping_patience == 0): for epoch in range(1, args.epochs): log = train_epoch() logs.append(log) # Validation metrics val_auc, val_ap = test(data.val_pos_edge_index, data.val_neg_edge_index) print('Validation-Epoch: {:03d}, AUC: {:.4f}, AP: {:.4f}'.format( epoch, val_auc, val_ap)) # Stop training if validation scores have not improved early_stopping(val_ap, model) if early_stopping.early_stop: print("Applying early-stopping") break else: epoch = 0 # Load best encoder print("Load best model for evaluation.") model.load_state_dict(torch.load('checkpoint.pt')) print( "__________________________________________________________________________" ) # Training is finished, calculate test metrics test_auc, test_ap = test(data.test_pos_edge_index, data.test_neg_edge_index) print('Test Results: {:03d}, AUC: {:.4f}, AP: {:.4f}'.format( epoch, test_auc, test_ap)) # Check if early stopping was applied or not - if not: model might not be done with training if args.epochs == epoch + 1: print("Model might need more epochs - Increase number of Epochs!") # Evaluate full graph latent_embeddings = model.encode(node_features, train_pos_edge_index) # Save embeddings to embeddings folder if flag is set if args.save_embeddings: embeddings_folder = osp.join(osp.dirname(osp.abspath(__file__)), 'embeddings') if not osp.isdir(embeddings_folder): os.makedirs(embeddings_folder) torch.save( latent_embeddings, osp.join(embeddings_folder, args.dataset + "_" + args.decoder + ".pt")) if not args.lsh: # Compute precision recall w.r.t the ground truth graph graph_precision, graph_recall = test_naive_graph(latent_embeddings) del model del encoder del decoder torch.cuda.empty_cache() else: # Precision w.r.t. the generated graph naive_precision, naive_recall, naive_time, naive_size, lsh_precision, \ lsh_recall, lsh_time, lsh_size, \ compare_precision, compare_recall = test_compare_lsh_naive_graphs( latent_embeddings) del model del encoder del decoder torch.cuda.empty_cache() return { 'args': args, 'test_auc': test_auc, 'test_ap': test_ap, 'naive_precision': naive_precision, 'naive_recall': naive_recall, 'naive_time': naive_time, 'naive_size': naive_size, 'lsh_precision': lsh_precision, 'lsh_recall': lsh_recall, 'lsh_time': lsh_time, 'lsh_size': lsh_size, 'compare_precision': compare_precision, 'compare_recall': compare_recall }
def run_GAE(input_data, output_dir, epochs=1000, lr=0.01, weight_decay=0.0005): device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print('Device: '.ljust(32), device) print('Model Name: '.ljust(32), 'GAE') print('Model params:{:19} lr: {} weight_decay: {}'.format( '', lr, weight_decay)) print('Total number of epochs to run: '.ljust(32), epochs) print('*' * 70) data = input_data.clone().to(device) in_channels = data.num_features out_channels = data.num_classes.item() model = GAE(GAEncoder(in_channels, out_channels)).to(device) data = input_data.clone().to(device) split_data = model.split_edges(data) x, train_pos_edge_index, edge_attr = split_data.x.to( device), split_data.train_pos_edge_index.to(device), data.edge_attr.to( device) split_data.train_idx = split_data.test_idx = data.y = None optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay) train_losses, test_losses = [], [] aucs = [] aps = [] model.train() for epoch in range(1, epochs + 1): train_loss = 0 test_loss = 0 optimizer.zero_grad() z = model.encode(x, train_pos_edge_index) train_loss = model.recon_loss(z, train_pos_edge_index) train_losses.append(train_loss) train_loss.backward() optimizer.step() model.eval() with torch.no_grad(): z = model.encode(x, train_pos_edge_index) auc, ap = model.test(z, split_data.test_pos_edge_index, split_data.test_neg_edge_index) test_loss = model.recon_loss(z, data.test_pos_edge_index) test_losses.append(test_loss.item()) aucs.append(auc) aps.append(ap) figname = os.path.join( output_dir, "_".join((GAE.__name__, str(lr), str(weight_decay)))) makepath(output_dir) if (epoch % int(epochs / 10) == 0): print( 'Epoch: {} Train loss: {} Test loss: {} AUC: {} AP: {}' .format(epoch, train_loss, test_loss, auc, ap)) if (epoch == epochs): print( '-' * 65, '\nFinal epoch: {} Train loss: {} Test loss: {} AUC: {} AP: {}' .format(epoch, train_loss, test_loss, auc, ap)) log = 'Final epoch: {} Train loss: {} Test loss: {} AUC: {} AP: {}'.format( epoch, train_loss, test_loss, auc, ap) write_log(log, figname) print('-' * 65) plot_linkpred(train_losses, test_losses, aucs, aps, output_dir, epochs, figname) return
def perturb_edges(data, name, remove_pct, add_pct, hidden_channels=16, epochs=400): if remove_pct == 0 and add_pct == 0: return try: cached = pickle.load( open(f'{ROOT}/cache/edge/{name}_{remove_pct}_{add_pct}.pt', 'rb')) print(f'Use cached edge augmentation for dataset {name}') if data.setting == 'inductive': data.train_edge_index = cached else: data.edge_index = cached return except FileNotFoundError: try: A_pred, adj_orig = pickle.load( open(f'{ROOT}/cache/edge/{name}.pt', 'rb')) A = sample_graph_det(adj_orig, A_pred, remove_pct, add_pct) data.edge_index, _ = from_scipy_sparse_matrix(A) pickle.dump( data.edge_index, open(f'{ROOT}/cache/edge/{name}_{remove_pct}_{add_pct}.pt', 'wb')) return except FileNotFoundError: print( f'cache/edge/{name}_{remove_pct}_{add_pct}.pt not found! Regenerating it now' ) device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') if data.setting == 'inductive': train_data = Data(x=data.train_x, ori_x=data.ori_x, edge_index=data.train_edge_index, y=data.train_y) else: train_data = deepcopy(data) edge_index = deepcopy(train_data.edge_index) train_data = train_test_split_edges(train_data, val_ratio=0.1, test_ratio=0) num_features = train_data.ori_x.shape[1] model = GAE(GCNEncoder(num_features, hidden_channels)) model = model.to(device) x = train_data.ori_x.to(device) train_pos_edge_index = train_data.train_pos_edge_index.to(device) optimizer = torch.optim.Adam(model.parameters(), lr=0.01) best_val_auc = 0 best_z = None for epoch in range(1, epochs + 1): model.train() optimizer.zero_grad() z = model.encode(x, train_pos_edge_index) loss = model.recon_loss(z, train_pos_edge_index) loss.backward() optimizer.step() model.eval() with torch.no_grad(): z = model.encode(x, train_pos_edge_index) auc, ap = model.test(z, train_data.val_pos_edge_index, train_data.val_neg_edge_index) print('Val | Epoch: {:03d}, AUC: {:.4f}, AP: {:.4f}'.format( epoch, auc, ap)) if auc > best_val_auc: best_val_auc = auc best_z = deepcopy(z) A_pred = torch.sigmoid(torch.mm(z, z.T)).cpu().numpy() adj_orig = to_scipy_sparse_matrix(edge_index).asformat('csr') adj_pred = sample_graph_det(adj_orig, A_pred, remove_pct, add_pct) if data.setting == 'inductive': data.train_edge_index, _ = from_scipy_sparse_matrix(adj_pred) else: data.edge_index, _ = from_scipy_sparse_matrix(adj_pred) pickle.dump((A_pred, adj_orig), open(f'{ROOT}/cache/edge/{name}.pt', 'wb')) if data.setting == 'inductive': pickle.dump( data.train_edge_index, open(f'{ROOT}/cache/edge/{name}_{remove_pct}_{add_pct}.pt', 'wb')) else: pickle.dump( data.edge_index, open(f'{ROOT}/cache/edge/{name}_{remove_pct}_{add_pct}.pt', 'wb'))