-
Notifications
You must be signed in to change notification settings - Fork 2
/
contrast_subgraph.py
29 lines (26 loc) · 956 Bytes
/
contrast_subgraph.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
import heapq
import math
import numpy as np
import torch
import dgl
from dgl.data.utils import load_graphs
from copy import deepcopy
from tqdm import tqdm
def get_summary_tensor(G_dataset, Labels, device, merge_classes=False):
num_G = len(G_dataset)
Labels = Labels.tolist()
node_num = G_dataset[0].ndata['feat'].shape[0]
adj_dict = {}
nodes_dict = {}
final_adj_dict = {}
final_nodes_dict = {}
for i in range(num_G):
if Labels[i] not in adj_dict.keys():
adj_dict[Labels[i]] = []
nodes_dict[Labels[i]] = []
adj_dict[Labels[i]].append(G_dataset[i].edata['feat'].squeeze().view(node_num, -1).tolist())
nodes_dict[Labels[i]].append(G_dataset[i].ndata['feat'].tolist())
for i in adj_dict.keys():
final_adj_dict[i] = torch.tensor(adj_dict[i]).to(device)
final_nodes_dict[i] = torch.tensor(nodes_dict[i]).to(device)
return final_adj_dict, final_nodes_dict