|
- import hetu
- import graphmix
- import numpy as np
- from tqdm import tqdm
-
-
- def padding(graph, target_num_nodes):
- assert graph.num_nodes <= target_num_nodes
- graph.convert2coo()
- new_graph = graphmix.Graph(graph.edge_index, target_num_nodes)
- new_graph.tag = graph.tag
- new_graph.type = graph.type
- extra = target_num_nodes - graph.num_nodes
- new_graph.i_feat = np.concatenate(
- [graph.i_feat, np.tile(graph.i_feat[0], [extra, 1])])
- new_graph.f_feat = np.concatenate(
- [graph.f_feat, np.tile(graph.f_feat[0], [extra, 1])])
- if graph.extra.size:
- new_graph.extra = np.concatenate([graph.extra, np.zeros([extra, 1])])
- return new_graph
-
-
- def prepare_data(ngraph):
- cli = graphmix.Client()
- graphs = []
- for i in tqdm(range(ngraph)):
- query = cli.pull_graph()
- graph = cli.wait(query)
- graphs.append(graph)
- max_num_nodes = 0
- for i in range(ngraph):
- max_num_nodes = max(max_num_nodes, graphs[i].num_nodes)
- for i in range(ngraph):
- graphs[i] = padding(graphs[i], max_num_nodes)
- return graphs
-
-
- def get_norm_adj(graph, device, use_original_gcn_norm=False):
- norm = graph.gcn_norm(use_original_gcn_norm)
- mp_mat = hetu.ndarray.sparse_array(
- values=norm,
- indices=(graph.edge_index[1], graph.edge_index[0]),
- shape=(graph.num_nodes, graph.num_nodes),
- ctx=device
- )
- return mp_mat
|
