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 load_data(dataset_name): if dataset_name in ['cora', 'citeseer', 'pubmed']: path = os.path.join(os.path.dirname(os.path.realpath(__file__)), '.', 'data', dataset_name) data = Planetoid(path, dataset_name)[0] else: data = load_wiki.load_data() data.edge_index = gutils.to_undirected(data.edge_index) data = GAE.split_edges(GAE, data) features = data.x.numpy() train_pos_edges = data.train_pos_edge_index.numpy() train_neg_edges = sample_negative(count=train_pos_edges.shape[1], avoid=train_pos_edges, nodes=features.shape[0]) x_tr, y_tr = combine_node_pair_features(features, train_pos_edges, train_neg_edges) x_val, y_val = combine_node_pair_features(features, data.val_pos_edge_index.numpy(), data.val_neg_edge_index.numpy()) x_test, y_test = combine_node_pair_features( features, data.test_pos_edge_index.numpy(), data.test_neg_edge_index.numpy()) return x_tr, y_tr, x_val, y_val, x_test, y_test
def load_data(dataset_name): path = osp.join(osp.dirname(osp.realpath(__file__)), '.', 'data', dataset_name) dataset = Planetoid(path, dataset_name, T.TargetIndegree()) num_features = dataset.num_features data = GAE.split_edges(GAE, dataset[0]) data.train_pos_edge_index = gutils.to_undirected(data.train_pos_edge_index) data.val_pos_edge_index = gutils.to_undirected(data.val_pos_edge_index) data.test_pos_edge_index = gutils.to_undirected(data.test_pos_edge_index) data.edge_index = torch.cat([ data.train_pos_edge_index, data.val_pos_edge_index, data.test_pos_edge_index ], dim=1) data.edge_train_mask = torch.cat([ torch.ones((data.train_pos_edge_index.size(-1))), torch.zeros((data.val_pos_edge_index.size(-1))), torch.zeros((data.test_pos_edge_index.size(-1))) ], dim=0).byte() data.edge_val_mask = torch.cat([ torch.zeros((data.train_pos_edge_index.size(-1))), torch.ones((data.val_pos_edge_index.size(-1))), torch.zeros((data.test_pos_edge_index.size(-1))) ], dim=0).byte() data.edge_test_mask = torch.cat([ torch.zeros((data.train_pos_edge_index.size(-1))), torch.zeros((data.val_pos_edge_index.size(-1))), torch.ones((data.test_pos_edge_index.size(-1))) ], dim=0).byte() data.edge_type = torch.zeros(((data.edge_index.size(-1)), )).long() data.batch = torch.zeros((1, data.num_nodes), dtype=torch.int64).view(-1) data.num_graphs = 1 return data, num_features
parser.add_argument('--dataset') parser.add_argument('--epochs', type=int, default=200) parser.add_argument('--val-freq', type=int, default=20) parser.add_argument('--runs', type=int, default=10) parser.add_argument('--test', action='store_true', default=False) args = parser.parse_args() if args.dataset in ['cora', 'citeseer', 'pubmed']: path = os.path.join(os.path.dirname(os.path.realpath(__file__)), '.', 'data', args.dataset) data = Planetoid(path, args.dataset)[0] else: data = load_wiki.load_data() data.edge_index = gutils.to_undirected(data.edge_index) data = GAE.split_edges(GAE, data) num_features = data.x.shape[1] aucs = [] aps = [] for run in range(args.runs): model = VGAE(VGAE_Encoder(num_features)) optimizer = torch.optim.Adam(model.parameters(), lr=0.01) # Training loop for epoch in range(args.epochs): model.train() optimizer.zero_grad() z = model.encode(data.x, data.train_pos_edge_index) loss = model.recon_loss( z, data.train_pos_edge_index) #0.01*model.kl_loss()
class Encoder(torch.nn.Module): def __init__(self, in_channels, out_channels): super(Encoder, self).__init__() self.conv1 = GCNConv(in_channels, 2 * out_channels) self.conv2 = GCNConv(2 * out_channels, out_channels) def forward(self, x, edge_index): x = F.relu(self.conv1(x, edge_index)) return self.conv2(x, edge_index) device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = GAE(Encoder(dataset.num_features, out_channels=16)).to(device) data.train_mask = data.val_mask = data.test_mask = data.y = None data = model.split_edges(data) x, edge_index = data.x.to(device), data.edge_index.to(device) optimizer = torch.optim.Adam(model.parameters(), lr=0.01) def train(): model.train() optimizer.zero_grad() z = model.encode(x, edge_index) loss = model.loss(z, data.train_pos_edge_index, data.train_neg_adj_mask) loss.backward() optimizer.step() def test(pos_edge_index, neg_edge_index): model.eval()
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