def test_set_batch_info(idtype): ctx = F.ctx() g1 = dgl.rand_graph(30, 100).astype(idtype).to(F.ctx()) g2 = dgl.rand_graph(40, 200).astype(idtype).to(F.ctx()) bg = dgl.batch([g1, g2]) batch_num_nodes = F.astype(bg.batch_num_nodes(), idtype) batch_num_edges = F.astype(bg.batch_num_edges(), idtype) # test homogeneous node subgraph sg_n = dgl.node_subgraph(bg, list(range(10, 20)) + list(range(50, 60))) induced_nodes = sg_n.ndata['_ID'] induced_edges = sg_n.edata['_ID'] new_batch_num_nodes = _get_subgraph_batch_info(bg.ntypes, [induced_nodes], batch_num_nodes) new_batch_num_edges = _get_subgraph_batch_info(bg.canonical_etypes, [induced_edges], batch_num_edges) sg_n.set_batch_num_nodes(new_batch_num_nodes) sg_n.set_batch_num_edges(new_batch_num_edges) subg_n1, subg_n2 = dgl.unbatch(sg_n) subg1 = dgl.node_subgraph(g1, list(range(10, 20))) subg2 = dgl.node_subgraph(g2, list(range(20, 30))) assert subg_n1.num_edges() == subg1.num_edges() assert subg_n2.num_edges() == subg2.num_edges() # test homogeneous edge subgraph sg_e = dgl.edge_subgraph(bg, list(range(40, 70)) + list(range(150, 200)), relabel_nodes=False) induced_nodes = F.arange(0, bg.num_nodes(), idtype) induced_edges = sg_e.edata['_ID'] new_batch_num_nodes = _get_subgraph_batch_info(bg.ntypes, [induced_nodes], batch_num_nodes) new_batch_num_edges = _get_subgraph_batch_info(bg.canonical_etypes, [induced_edges], batch_num_edges) sg_e.set_batch_num_nodes(new_batch_num_nodes) sg_e.set_batch_num_edges(new_batch_num_edges) subg_e1, subg_e2 = dgl.unbatch(sg_e) subg1 = dgl.edge_subgraph(g1, list(range(40, 70)), relabel_nodes=False) subg2 = dgl.edge_subgraph(g2, list(range(50, 100)), relabel_nodes=False) assert subg_e1.num_nodes() == subg1.num_nodes() assert subg_e2.num_nodes() == subg2.num_nodes()
def prepare_data_for_year(graph, features, labels, years, current_year, history, exclude_class=None, device=None, backend='dgl', num_hops=None): print("Preparing data for year", current_year) # Prepare subgraph subg_node_mask = ((years <= current_year) & (years >= (current_year - history))) subg_nodes = torch.arange(features.size(0))[subg_node_mask] subg_num_nodes = subg_nodes.size(0) if backend == 'dgl': print("Creating dgl subgraph") subg = dgl.node_subgraph(graph, subg_nodes) print("Subgraph type:", type(subg)) subg.set_n_initializer(dgl.init.zero_initializer) elif backend == 'geometric': print("Creating geometric subgraph") subg, __edge_attr = tg.utils.subgraph(subg_node_mask, graph, relabel_nodes=True, num_nodes=subg_num_nodes) else: raise ValueError("Unkown backend: " + backend) subg_features = features[subg_nodes] subg_labels = labels[subg_nodes] subg_years = years[subg_nodes] # Prepare masks wrt *subgraph* # train_nid = torch.arange(subg_num_nodes)[subg_years < current_year] # test_nid = torch.arange(subg_num_nodes)[subg_years == current_year] # print("[{}] #Training: {}".format(current_year, train_nid.size(0))) # print("[{}] #Test : {}".format(current_year, test_nid.size(0))) train_nid = subg_years < current_year test_nid = subg_years == current_year print("[{}] #Training: {}".format(current_year, train_nid.sum())) print("[{}] #Test : {}".format(current_year, test_nid.sum())) if device is not None: subg = subg.to(device) subg_features = subg_features.to(device) subg_labels = subg_labels.to(device) # train_nid = train_nid.to(device) # test_nid = test_nid.to(device) return subg, subg_features, subg_labels, subg_years, train_nid, test_nid
def sample_subgraph(self, target_nodes): """ Args: target_nodes(Tensor): Tensor of two target nodes Returns: subgraph(DGLGraph): subgraph """ sample_nodes = [target_nodes] frontiers = target_nodes for i in range(self.hop): frontiers = self.graph.out_edges(frontiers)[1] frontiers = torch.unique(frontiers) sample_nodes.append(frontiers) sample_nodes = torch.cat(sample_nodes) sample_nodes = torch.unique(sample_nodes) subgraph = dgl.node_subgraph(self.graph, sample_nodes) # Each node should have unique node id in the new subgraph u_id = int( torch.nonzero(subgraph.ndata[NID] == int(target_nodes[0]), as_tuple=False)) v_id = int( torch.nonzero(subgraph.ndata[NID] == int(target_nodes[1]), as_tuple=False)) # remove link between target nodes in positive subgraphs. # Edge removing will rearange NID and EID, which lose the original NID and EID. # if dgl.__version__[:5] < '0.6.0': # nids = subgraph.ndata[NID] # eids = subgraph.edata[EID] # if subgraph.has_edges_between(u_id, v_id): # link_id = subgraph.edge_ids(u_id, v_id, return_uv=True)[2] # subgraph.remove_edges(link_id) # eids = eids[subgraph.edata[EID]] # if subgraph.has_edges_between(v_id, u_id): # link_id = subgraph.edge_ids(v_id, u_id, return_uv=True)[2] # subgraph.remove_edges(link_id) # eids = eids[subgraph.edata[EID]] # subgraph.ndata[NID] = nids # subgraph.edata[EID] = eids if subgraph.has_edges_between(u_id, v_id): link_id = subgraph.edge_ids(u_id, v_id, return_uv=True)[2] subgraph.remove_edges(link_id) if subgraph.has_edges_between(v_id, u_id): link_id = subgraph.edge_ids(v_id, u_id, return_uv=True)[2] subgraph.remove_edges(link_id) z = drnl_node_labeling(subgraph, u_id, v_id) subgraph.ndata['z'] = z return subgraph
def sample_subgraph_dgl(whole_data): data = whole_data[0] # dgl data # find data with different labels # random walk start = [ random.randint(0, data.num_nodes - 1) for i in range(self.subgraphs) ] traces, _ = dgl.sampling.random_walk_with_restart( data, start, length=self.sample_batch_size, restart_prob=1 / self.sample_walk_length) subgraphs = dgl.node_subgraph( data, [traces[i, :] for i in traces.size(0)]) return subgraphs
def forward(self, graph: dgl.DGLGraph, feature: torch.Tensor): score = self.score_layer(graph, feature).squeeze() perm, next_batch_num_nodes = topk( score, self.ratio, get_batch_id(graph.batch_num_nodes()), graph.batch_num_nodes()) feature = feature[perm] * self.non_linearity(score[perm]).view(-1, 1) graph = dgl.node_subgraph(graph, perm) # node_subgraph currently does not support batch-graph, # the 'batch_num_nodes' of the result subgraph is None. # So we manually set the 'batch_num_nodes' here. # Since global pooling has nothing to do with 'batch_num_edges', # we can leave it to be None or unchanged. graph.set_batch_num_nodes(next_batch_num_nodes) return graph, feature, perm
def forward(self, graph: DGLGraph, feat: Tensor, select_idx: Tensor, non_select_idx: Optional[Tensor] = None, scores: Optional[Tensor] = None, pool_graph=False): """ Description ----------- Perform graph pooling. Parameters ---------- graph : dgl.DGLGraph The input graph feat : torch.Tensor The input node feature select_idx : torch.Tensor The index in fine graph of node from coarse graph, this is obtained from previous graph pooling layers. non_select_idx : torch.Tensor, optional The index that not included in output graph. default: :obj:`None` scores : torch.Tensor, optional Scores for nodes used for pooling and scaling. default: :obj:`None` pool_graph : bool, optional Whether perform graph pooling on graph topology. default: :obj:`False` """ if self.use_gcn: feat = self.down_sample_gcn(graph, feat) feat = feat[select_idx] if scores is not None: feat = feat * scores.unsqueeze(-1) if pool_graph: num_node_batch = graph.batch_num_nodes() graph = dgl.node_subgraph(graph, select_idx) graph.set_batch_num_nodes(num_node_batch) return feat, graph else: return feat
def extract_graph_new(self, G, u_id, v_id): v_id += self.num_user static_u = torch.zeros(len(self.class_values)) static_v = torch.zeros(len(self.class_values)) start0 = time.time() u_nodes, v, e_ids_1 = G.in_edges(v_id, "all") u, v_nodes, e_ids_2 = G.out_edges(u_id, "all") e_ids = [] nodes = torch.cat([u_nodes, v_nodes]) for i in range(u_nodes.shape[0]): if u_nodes[i] == u_id: e_ids.append(e_ids_1[i]) for i in range(v_nodes.shape[0]): if v_nodes[i] == v_id: e_ids.append(e_ids_2[i]) #start1 = time.time() #print(start1 - start0) subg = dgl.node_subgraph(G, nodes) #start2 = time.time() #print(start2 - start1) subg.ndata['node_label'] = torch.zeros([subg.num_nodes(), 4]) pid = subg.ndata[dgl.NID] #start3 = time.time() #print(start3 - start2) for i in range(pid.shape[0]): if pid[i] == u_id: e_u = i subg.ndata['node_label'][i, 0] = 1 elif pid[i] == v_id: e_v = i subg.ndata['node_label'][i, 1] = 1 elif pid[i] in u: subg.ndata['node_label'][i, 2] = 1 elif pid[i] in v: subg.ndata['node_label'][i, 3] = 1 subg = dgl.remove_edges(subg, e_ids) start6 = time.time() print(start6 - start0) print() return subg
def sample_subgraph(self, target_nodes): """ Args: target_nodes(Tensor): Tensor of two target nodes Returns: subgraph(DGLGraph): subgraph """ sample_nodes = [target_nodes] frontiers = target_nodes for i in range(self.hop): frontiers = self.graph.out_edges(frontiers)[1] frontiers = torch.unique(frontiers) sample_nodes.append(frontiers) sample_nodes = torch.cat(sample_nodes) sample_nodes = torch.unique(sample_nodes) subgraph = dgl.node_subgraph(self.graph, sample_nodes) # Each node should have unique node id in the new subgraph u_id = int( torch.nonzero(subgraph.ndata[NID] == int(target_nodes[0]), as_tuple=False)) v_id = int( torch.nonzero(subgraph.ndata[NID] == int(target_nodes[1]), as_tuple=False)) # remove link between target nodes in positive subgraphs. if subgraph.has_edges_between(u_id, v_id): link_id = subgraph.edge_ids(u_id, v_id, return_uv=True)[2] subgraph.remove_edges(link_id) if subgraph.has_edges_between(v_id, u_id): link_id = subgraph.edge_ids(v_id, u_id, return_uv=True)[2] subgraph.remove_edges(link_id) z = drnl_node_labeling(subgraph, u_id, v_id) subgraph.ndata['z'] = z return subgraph
def step(self, pick_nodes, kick_nodes=[]): self.picked_nodes = list(set(self.picked_nodes + pick_nodes.tolist())) self.picked_nodes = [ x for x in self.picked_nodes if x not in kick_nodes ] self.graph = dgl.node_subgraph(self.dataset, self.picked_nodes) features = torch.stack([ x for i, x in enumerate(self.dataset.ndata['feat']) if i in self.picked_nodes ]) labels = torch.stack([ x for i, x in enumerate(self.dataset.ndata['label']) if i in self.picked_nodes ]) loss_fcn = torch.nn.CrossEntropyLoss() logits = self.GNN(self.graph, features) loss = loss_fcn(logits, labels) if len(self.picked_nodes) >= 2000: self.done = True return -loss, self.done
lr=args.lr1, weight_decay=args.wd1) loss_fn = nn.BCEWithLogitsLoss() node_list = list(range(n_node)) # Step 4: Training epochs ================================================================ # best = float('inf') cnt_wait = 0 for epoch in range(args.epochs): model.train() optimizer.zero_grad() sample_idx = random.sample(node_list, sample_size) g = dgl.node_subgraph(graph, sample_idx) dg = dgl.node_subgraph(diff_graph, sample_idx) f = g.ndata.pop('feat') ew = dg.edata.pop('edge_weight') shuf_idx = np.random.permutation(sample_size) sf = f[shuf_idx, :] g = g.to(args.device) dg = dg.to(args.device) f = f.to(args.device) ew = ew.to(args.device) sf = sf.to(args.device)
def train(args, device): elliptic_dataset = EllipticDataset(raw_dir=args.raw_dir, processed_dir=args.processed_dir, self_loop=True, reverse_edge=True) g, node_mask_by_time = elliptic_dataset.process() num_classes = elliptic_dataset.num_classes cached_subgraph = [] cached_labeled_node_mask = [] for i in range(len(node_mask_by_time)): # we add self loop edge when we construct full graph, not here node_subgraph = dgl.node_subgraph(graph=g, nodes=node_mask_by_time[i]) cached_subgraph.append(node_subgraph.to(device)) valid_node_mask = node_subgraph.ndata['label'] >= 0 cached_labeled_node_mask.append(valid_node_mask) if args.model == 'EvolveGCN-O': model = EvolveGCNO(in_feats=int(g.ndata['feat'].shape[1]), n_hidden=args.n_hidden, num_layers=args.n_layers) elif args.model == 'EvolveGCN-H': model = EvolveGCNH(in_feats=int(g.ndata['feat'].shape[1]), num_layers=args.n_layers) else: return NotImplementedError('Unsupported model {}'.format(args.model)) model = model.to(device) optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) # split train, valid, test(0-30,31-35,36-48) # train/valid/test split follow the paper. train_max_index = 30 valid_max_index = 35 test_max_index = 48 time_window_size = args.n_hist_steps loss_class_weight = [float(w) for w in args.loss_class_weight.split(',')] loss_class_weight = torch.Tensor(loss_class_weight).to(device) train_measure = Measure(num_classes=num_classes, target_class=args.eval_class_id) valid_measure = Measure(num_classes=num_classes, target_class=args.eval_class_id) test_measure = Measure(num_classes=num_classes, target_class=args.eval_class_id) test_res_f1 = 0 for epoch in range(args.num_epochs): model.train() for i in range(time_window_size, train_max_index + 1): g_list = cached_subgraph[i - time_window_size:i + 1] predictions = model(g_list) # get predictions which has label predictions = predictions[cached_labeled_node_mask[i]] labels = cached_subgraph[i].ndata['label'][ cached_labeled_node_mask[i]].long() loss = F.cross_entropy(predictions, labels, weight=loss_class_weight) optimizer.zero_grad() loss.backward() optimizer.step() train_measure.append_measures(predictions, labels) # get each epoch measures during training. cl_precision, cl_recall, cl_f1 = train_measure.get_total_measure() train_measure.update_best_f1(cl_f1, epoch) # reset measures for next epoch train_measure.reset_info() print( "Train Epoch {} | class {} | precision:{:.4f} | recall: {:.4f} | f1: {:.4f}" .format(epoch, args.eval_class_id, cl_precision, cl_recall, cl_f1)) # eval model.eval() for i in range(train_max_index + 1, valid_max_index + 1): g_list = cached_subgraph[i - time_window_size:i + 1] predictions = model(g_list) # get node predictions which has label predictions = predictions[cached_labeled_node_mask[i]] labels = cached_subgraph[i].ndata['label'][ cached_labeled_node_mask[i]].long() valid_measure.append_measures(predictions, labels) # get each epoch measure during eval. cl_precision, cl_recall, cl_f1 = valid_measure.get_total_measure() valid_measure.update_best_f1(cl_f1, epoch) # reset measures for next epoch valid_measure.reset_info() print( "Eval Epoch {} | class {} | precision:{:.4f} | recall: {:.4f} | f1: {:.4f}" .format(epoch, args.eval_class_id, cl_precision, cl_recall, cl_f1)) # early stop if epoch - valid_measure.target_best_f1_epoch >= args.patience: print("Best eval Epoch {}, Cur Epoch {}".format( valid_measure.target_best_f1_epoch, epoch)) break # if cur valid f1 score is best, do test if epoch == valid_measure.target_best_f1_epoch: print("###################Epoch {} Test###################".format( epoch)) for i in range(valid_max_index + 1, test_max_index + 1): g_list = cached_subgraph[i - time_window_size:i + 1] predictions = model(g_list) # get predictions which has label predictions = predictions[cached_labeled_node_mask[i]] labels = cached_subgraph[i].ndata['label'][ cached_labeled_node_mask[i]].long() test_measure.append_measures(predictions, labels) # we get each subgraph measure when testing to match fig 4 in EvolveGCN paper. cl_precisions, cl_recalls, cl_f1s = test_measure.get_each_timestamp_measure( ) for index, (sub_p, sub_r, sub_f1) in enumerate( zip(cl_precisions, cl_recalls, cl_f1s)): print( " Test | Time {} | precision:{:.4f} | recall: {:.4f} | f1: {:.4f}" .format(valid_max_index + index + 2, sub_p, sub_r, sub_f1)) # get each epoch measure during test. cl_precision, cl_recall, cl_f1 = test_measure.get_total_measure() test_measure.update_best_f1(cl_f1, epoch) # reset measures for next test test_measure.reset_info() test_res_f1 = cl_f1 print( " Test | Epoch {} | class {} | precision:{:.4f} | recall: {:.4f} | f1: {:.4f}" .format(epoch, args.eval_class_id, cl_precision, cl_recall, cl_f1)) print("Best test f1 is {}, in Epoch {}".format( test_measure.target_best_f1, test_measure.target_best_f1_epoch)) if test_measure.target_best_f1_epoch != valid_measure.target_best_f1_epoch: print( "The Epoch get best Valid measure not get the best Test measure, " "please checkout the test result in Epoch {}, which f1 is {}". format(valid_measure.target_best_f1_epoch, test_res_f1))
def main(args): torch.manual_seed(args.rnd_seed) np.random.seed(args.rnd_seed) random.seed(args.rnd_seed) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False multitask_data = set(['ppi']) multitask = args.dataset in multitask_data # load and preprocess dataset data = load_data(args) g = data.g train_mask = g.ndata['train_mask'] val_mask = g.ndata['val_mask'] test_mask = g.ndata['test_mask'] labels = g.ndata['label'] train_nid = np.nonzero(train_mask.data.numpy())[0].astype(np.int64) # Normalize features if args.normalize: feats = g.ndata['feat'] train_feats = feats[train_mask] scaler = sklearn.preprocessing.StandardScaler() scaler.fit(train_feats.data.numpy()) features = scaler.transform(feats.data.numpy()) g.ndata['feat'] = torch.FloatTensor(features) in_feats = g.ndata['feat'].shape[1] n_classes = data.num_classes n_edges = g.number_of_edges() n_train_samples = train_mask.int().sum().item() n_val_samples = val_mask.int().sum().item() n_test_samples = test_mask.int().sum().item() print("""----Data statistics------' #Edges %d #Classes %d #Train samples %d #Val samples %d #Test samples %d""" % (n_edges, n_classes, n_train_samples, n_val_samples, n_test_samples)) # create GCN model if args.self_loop and not args.dataset.startswith('reddit'): g = dgl.remove_self_loop(g) g = dgl.add_self_loop(g) print("adding self-loop edges") # metis only support int64 graph g = g.long() if args.use_pp: g.update_all(fn.copy_u('feat', 'm'), fn.sum('m', 'feat_agg')) g.ndata['feat'] = torch.cat([g.ndata['feat'], g.ndata['feat_agg']], 1) del g.ndata['feat_agg'] cluster_iterator = dgl.dataloading.GraphDataLoader( dgl.dataloading.ClusterGCNSubgraphIterator( dgl.node_subgraph(g, train_nid), args.psize, './cache'), batch_size=args.batch_size, num_workers=4) #cluster_iterator = ClusterIter( # args.dataset, g, args.psize, args.batch_size, train_nid, use_pp=args.use_pp) # set device for dataset tensors if args.gpu < 0: cuda = False else: cuda = True torch.cuda.set_device(args.gpu) val_mask = val_mask.cuda() test_mask = test_mask.cuda() g = g.int().to(args.gpu) print('labels shape:', g.ndata['label'].shape) print("features shape, ", g.ndata['feat'].shape) model = GraphSAGE(in_feats, args.n_hidden, n_classes, args.n_layers, F.relu, args.dropout, args.use_pp) if cuda: model.cuda() # logger and so on log_dir = save_log_dir(args) logger = Logger(os.path.join(log_dir, 'loggings')) logger.write(args) # Loss function if multitask: print('Using multi-label loss') loss_f = nn.BCEWithLogitsLoss() else: print('Using multi-class loss') loss_f = nn.CrossEntropyLoss() # use optimizer optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) # set train_nids to cuda tensor if cuda: train_nid = torch.from_numpy(train_nid).cuda() print("current memory after model before training", torch.cuda.memory_allocated(device=train_nid.device) / 1024 / 1024) start_time = time.time() best_f1 = -1 for epoch in range(args.n_epochs): for j, cluster in enumerate(cluster_iterator): # sync with upper level training graph if cuda: cluster = cluster.to(torch.cuda.current_device()) model.train() # forward batch_labels = cluster.ndata['label'] batch_train_mask = cluster.ndata['train_mask'] if batch_train_mask.sum().item() == 0: continue pred = model(cluster) loss = loss_f(pred[batch_train_mask], batch_labels[batch_train_mask]) optimizer.zero_grad() loss.backward() optimizer.step() # in PPI case, `log_every` is chosen to log one time per epoch. # Choose your log freq dynamically when you want more info within one epoch if j % args.log_every == 0: print(f"epoch:{epoch}/{args.n_epochs}, Iteration {j}/" f"{len(cluster_iterator)}:training loss", loss.item()) print("current memory:", torch.cuda.memory_allocated(device=pred.device) / 1024 / 1024) # evaluate if epoch % args.val_every == 0: val_f1_mic, val_f1_mac = evaluate( model, g, labels, val_mask, multitask) print( "Val F1-mic{:.4f}, Val F1-mac{:.4f}". format(val_f1_mic, val_f1_mac)) if val_f1_mic > best_f1: best_f1 = val_f1_mic print('new best val f1:', best_f1) torch.save(model.state_dict(), os.path.join( log_dir, 'best_model.pkl')) end_time = time.time() print(f'training using time {start_time-end_time}') # test if args.use_val: model.load_state_dict(torch.load(os.path.join( log_dir, 'best_model.pkl'))) test_f1_mic, test_f1_mac = evaluate( model, g, labels, test_mask, multitask) print("Test F1-mic{:.4f}, Test F1-mac{:.4f}". format(test_f1_mic, test_f1_mac))
compact_g.ndata['orig_id'] = th.as_tensor(per_type_ids) compact_g.ndata[dgl.NTYPE] = th.as_tensor(ntype) compact_g.ndata[dgl.NID] = th.as_tensor(uniq_ids) compact_g.ndata['inner_node'] = th.as_tensor(np.logical_and( uniq_ids >= nid_range[0], uniq_ids <= nid_range[1])) local_nids = compact_g.ndata[dgl.NID][compact_g.ndata['inner_node'].bool()] assert np.all((local_nids == th.arange( local_nids[0], local_nids[-1] + 1)).numpy()) print('|V|={}'.format(compact_g.number_of_nodes())) print('|E|={}'.format(compact_g.number_of_edges())) # We need to reshuffle nodes in a partition so that all local nodes are labelled starting from 0. reshuffle_nodes = th.arange(compact_g.number_of_nodes()) reshuffle_nodes = th.cat([reshuffle_nodes[compact_g.ndata['inner_node'].bool()], reshuffle_nodes[compact_g.ndata['inner_node'] == 0]]) compact_g1 = dgl.node_subgraph(compact_g, reshuffle_nodes) compact_g1.ndata['orig_id'] = compact_g.ndata['orig_id'][reshuffle_nodes] compact_g1.ndata[dgl.NTYPE] = compact_g.ndata[dgl.NTYPE][reshuffle_nodes] compact_g1.ndata[dgl.NID] = compact_g.ndata[dgl.NID][reshuffle_nodes] compact_g1.ndata['inner_node'] = compact_g.ndata['inner_node'][reshuffle_nodes] compact_g1.edata['orig_id'] = compact_g.edata['orig_id'][compact_g1.edata[dgl.EID]] compact_g1.edata[dgl.ETYPE] = compact_g.edata[dgl.ETYPE][compact_g1.edata[dgl.EID]] compact_g1.edata['inner_edge'] = compact_g.edata['inner_edge'][compact_g1.edata[dgl.EID]] # reshuffle edges on ETYPE as node_subgraph relabels edges idx = th.argsort(compact_g1.edata[dgl.ETYPE]) u, v = compact_g1.edges() u = u[idx] v = v[idx] compact_g2 = dgl.graph((u, v)) compact_g2.ndata['orig_id'] = compact_g1.ndata['orig_id']
def forward(self, graph: DGLGraph, feat: Tensor, e_feat=None): # top-k pool first if e_feat is None: e_feat = torch.ones((graph.number_of_edges(), ), dtype=feat.dtype, device=feat.device) batch_num_nodes = graph.batch_num_nodes() x_score = self.calc_info_score(graph, feat, e_feat) perm, next_batch_num_nodes = topk(x_score, self.ratio, get_batch_id(batch_num_nodes), batch_num_nodes) feat = feat[perm] pool_graph = None if not self.sample or not self.sl: # pool graph graph.edata["e"] = e_feat pool_graph = dgl.node_subgraph(graph, perm) e_feat = pool_graph.edata.pop("e") pool_graph.set_batch_num_nodes(next_batch_num_nodes) # no structure learning layer, directly return. if not self.sl: return pool_graph, feat, e_feat, perm # Structure Learning if self.sample: # A fast mode for large graphs. # In large graphs, learning the possible edge weights between each # pair of nodes is time consuming. To accelerate this process, # we sample it's K-Hop neighbors for each node and then learn the # edge weights between them. # first build multi-hop graph row, col = graph.all_edges() num_nodes = graph.num_nodes() scipy_adj = scipy.sparse.coo_matrix( (e_feat.detach().cpu(), (row.detach().cpu(), col.detach().cpu())), shape=(num_nodes, num_nodes)) for _ in range(self.k_hop): two_hop = scipy_adj**2 two_hop = two_hop * (1e-5 / two_hop.max()) scipy_adj = two_hop + scipy_adj row, col = scipy_adj.nonzero() row = torch.tensor(row, dtype=torch.long, device=graph.device) col = torch.tensor(col, dtype=torch.long, device=graph.device) e_feat = torch.tensor(scipy_adj.data, dtype=torch.float, device=feat.device) # perform pooling on multi-hop graph mask = perm.new_full((num_nodes, ), -1) i = torch.arange(perm.size(0), dtype=torch.long, device=perm.device) mask[perm] = i row, col = mask[row], mask[col] mask = (row >= 0) & (col >= 0) row, col = row[mask], col[mask] e_feat = e_feat[mask] # add remaining self loops mask = row != col num_nodes = perm.size(0) # num nodes after pool loop_index = torch.arange(0, num_nodes, dtype=row.dtype, device=row.device) inv_mask = ~mask loop_weight = torch.full((num_nodes, ), 0, dtype=e_feat.dtype, device=e_feat.device) remaining_e_feat = e_feat[inv_mask] if remaining_e_feat.numel() > 0: loop_weight[row[inv_mask]] = remaining_e_feat e_feat = torch.cat([e_feat[mask], loop_weight], dim=0) row, col = row[mask], col[mask] row = torch.cat([row, loop_index], dim=0) col = torch.cat([col, loop_index], dim=0) # attention scores weights = (torch.cat([feat[row], feat[col]], dim=1) * self.att).sum(dim=-1) weights = F.leaky_relu(weights, self.negative_slop) + e_feat * self.lamb # sl and normalization sl_graph = dgl.graph((row, col)) if self.sparse: weights = edge_sparsemax(sl_graph, weights) else: weights = edge_softmax(sl_graph, weights) # get final graph mask = torch.abs(weights) > 0 row, col, weights = row[mask], col[mask], weights[mask] pool_graph = dgl.graph((row, col)) pool_graph.set_batch_num_nodes(next_batch_num_nodes) e_feat = weights else: # Learning the possible edge weights between each pair of # nodes in the pooled subgraph, relative slower. # construct complete graphs for all graph in the batch # use dense to build, then transform to sparse. # maybe there's more efficient way? batch_num_nodes = next_batch_num_nodes block_begin_idx = torch.cat([ batch_num_nodes.new_zeros(1), batch_num_nodes.cumsum(dim=0)[:-1] ], dim=0) block_end_idx = batch_num_nodes.cumsum(dim=0) dense_adj = torch.zeros( (pool_graph.num_nodes(), pool_graph.num_nodes()), dtype=torch.float, device=feat.device) for idx_b, idx_e in zip(block_begin_idx, block_end_idx): dense_adj[idx_b:idx_e, idx_b:idx_e] = 1. row, col = torch.nonzero(dense_adj).t().contiguous() # compute weights for node-pairs weights = (torch.cat([feat[row], feat[col]], dim=1) * self.att).sum(dim=-1) weights = F.leaky_relu(weights, self.negative_slop) dense_adj[row, col] = weights # add pooled graph structure to weight matrix pool_row, pool_col = pool_graph.all_edges() dense_adj[pool_row, pool_col] += self.lamb * e_feat weights = dense_adj[row, col] del dense_adj torch.cuda.empty_cache() # edge softmax/sparsemax complete_graph = dgl.graph((row, col)) if self.sparse: weights = edge_sparsemax(complete_graph, weights) else: weights = edge_softmax(complete_graph, weights) # get new e_feat and graph structure, clean up. mask = torch.abs(weights) > 1e-9 row, col, weights = row[mask], col[mask], weights[mask] e_feat = weights pool_graph = dgl.graph((row, col)) pool_graph.set_batch_num_nodes(next_batch_num_nodes) return pool_graph, feat, e_feat, perm
import numpy as np import networkx as nx from karateclub.node_embedding.attributed import TENE import dgl from dgl.data import CoraGraphDataset data = CoraGraphDataset() g = data[0] g = dgl.add_self_loop( dgl.node_subgraph(g, list(set(np.random.choice(len(g.nodes()), 5))))) X = np.array(g.ndata['feat']) g = dgl.to_networkx(g).to_undirected() # g = nx.newman_watts_strogatz_graph(200, 20, 0.05) # # X = np.random.uniform(0, 1, (200, 200)) model = TENE() model.fit(g, X) embedding = model.get_embedding() embedding = np.sum(embedding, axis=0) # print(np.concatenate((embedding, embedding), axis=0).shape) print(embedding.shape)
def reset(self): self.picked_nodes = list( set(np.random.choice(dataset.nodes().numpy(), 1))) self.graph = dgl.node_subgraph(self.dataset, self.picked_nodes)
import argparse import dgl from dgl.data import CoraGraphDataset, CiteseerGraphDataset, PubmedGraphDataset import networkx as nx import torch import torch.nn.functional as F import numpy as np from GAT.model import GAT data = CoraGraphDataset() dataset = data[0] picked_nodes = list(set(np.random.choice(dataset.nodes().numpy(), 200))) graph = dgl.node_subgraph(dataset, picked_nodes) graph = dgl.add_self_loop(graph) features = torch.stack( [x for i, x in enumerate(dataset.ndata['feat']) if i in picked_nodes]) g = dgl.to_networkx(graph) f = np.array(dataset.ndata['feat']) print("*************") print(graph.edata['_ID']) parser = argparse.ArgumentParser(description='GAT') parser.add_argument("--gpu", type=int, default=0, help="which GPU to use. Set -1 to use CPU.") parser.add_argument("--epochs", type=int, default=500, help="number of training epochs") parser.add_argument("--num-heads",