Merge pull request #98 from ZHUI/graphsaint

add graphsaint support

ogb_examples/nodeproppred/ogbn-arxiv/dataloader/ogbn_arxiv_dataloader.py

@@ -24,7 +24,7 @@ import ssl
 ssl._create_default_https_context = ssl._create_unverified_context
 
 from pgl.contrib.ogb.nodeproppred.dataset_pgl import PglNodePropPredDataset
-#from pgl.sample import graph_saint_random_walk_sample
+from pgl.sample import graph_saint_random_walk_sample
 from ogb.nodeproppred import Evaluator
 import tqdm
 from collections import namedtuple
@@ -78,10 +78,10 @@ def k_hop_sampler(graph, samples, batch_nodes):
     return subgraph, sub_node_index
 
 
-#def graph_saint_randomwalk_sampler(graph, batch_nodes, max_depth=3):
-#    subgraph = graph_saint_random_walk_sample(graph, batch_nodes, max_depth)
-#    sub_node_index = subgraph.reindex_from_parrent_nodes(batch_nodes)
-#    return subgraph, sub_node_index
+def graph_saint_randomwalk_sampler(graph, batch_nodes, max_depth=3):
+    subgraph = graph_saint_random_walk_sample(graph, batch_nodes, max_depth)
+    sub_node_index = subgraph.reindex_from_parrent_nodes(batch_nodes)
+    return subgraph, sub_node_index
 
 
 class ArxivDataGenerator(BaseDataGenerator):

pgl/graph_kernel.pyx

@@ -321,3 +321,43 @@ def alias_sample_build_table(np.ndarray[np.float64_t, ndim=1] probs):
             if alias[l_i] < 1:
                 smaller_num.push_back(l_i)
     return alias, events
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def extract_edges_from_nodes(
+    np.ndarray[np.int64_t, ndim=1] adj_indptr,
+    np.ndarray[np.int64_t, ndim=1] sorted_v,
+    np.ndarray[np.int64_t, ndim=1] sorted_eid,
+    vector[long long] sampled_nodes,
+):
+    """
+    Extract all eids of given sampled_nodes for the origin graph.
+    ret_edge_index: edge ids between sampled_nodes.
+
+    Refers: https://github.com/GraphSAINT/GraphSAINT
+    """
+    cdef long long i, v, j
+    cdef long long num_v_orig, num_v_sub
+    cdef long long start_neigh, end_neigh
+    cdef vector[int] _arr_bit
+    cdef vector[long long] ret_edge_index
+    num_v_orig = adj_indptr.size-1
+    _arr_bit = vector[int](num_v_orig,-1)
+    num_v_sub = sampled_nodes.size()
+    i = 0
+    with nogil:
+        while i < num_v_sub:
+            _arr_bit[sampled_nodes[i]] = i
+            i = i + 1
+        i = 0
+        while i < num_v_sub:
+            v = sampled_nodes[i]
+            start_neigh = adj_indptr[v]
+            end_neigh = adj_indptr[v+1]
+            j = start_neigh
+            while j < end_neigh:
+                if _arr_bit[sorted_v[j]] > -1:
+                    ret_edge_index.push_back(sorted_eid[j])
+                j = j + 1
+            i = i + 1
+    return ret_edge_index

pgl/sample.py

@@ -24,7 +24,7 @@ from pgl import graph_kernel
 
 __all__ = [
     'graphsage_sample', 'node2vec_sample', 'deepwalk_sample',
-    'metapath_randomwalk', 'pinsage_sample'
+    'metapath_randomwalk', 'pinsage_sample', 'graph_saint_random_walk_sample'
 ]
 
 
@@ -55,7 +55,7 @@ def edge_hash(src, dst):
 
 def graphsage_sample(graph, nodes, samples, ignore_edges=[]):
     """Implement of graphsage sample.
-    
+
     Reference paper: https://cs.stanford.edu/people/jure/pubs/graphsage-nips17.pdf.
 
     Args:
@@ -63,7 +63,7 @@ def graphsage_sample(graph, nodes, samples, ignore_edges=[]):
         nodes: Sample starting from nodes
         samples: A list, number of neighbors in each layer
         ignore_edges: list of edge(src, dst) will be ignored.
-    
+
     Return:
         A list of subgraphs
     """
@@ -129,7 +129,7 @@ def alias_sample(size, alias, events):
         size: Output shape.
         alias: The alias table build by `alias_sample_build_table`.
         events: The events table build by `alias_sample_build_table`.
-    
+
     Return:
         samples: The generated random samples.
     """
@@ -283,13 +283,13 @@ def metapath_randomwalk(graph,
     Args:
         graph: instance of pgl heterogeneous graph
         start_nodes: start nodes to generate walk
-        metapath: meta path for sample nodes. 
+        metapath: meta path for sample nodes.
             e.g: "c2p-p2a-a2p-p2c"
         walk_length: the walk length
 
     Return:
-        a list of metapath walks. 
-        
+        a list of metapath walks.
+
     """
 
     edge_types = metapath.split('-')
@@ -390,18 +390,18 @@ def pinsage_sample(graph,
                    norm_bais=1.0,
                    ignore_edges=set()):
     """Implement of graphsage sample.
-    
+
     Reference paper: .
 
     Args:
         graph: A pgl graph instance
         nodes: Sample starting from nodes
         samples: A list, number of neighbors in each layer
-        top_k: select the top_k visit count nodes to construct the edges 
-        proba: the probability to return the origin node 
+        top_k: select the top_k visit count nodes to construct the edges
+        proba: the probability to return the origin node
         norm_bais: the normlization for the visit count
         ignore_edges: list of edge(src, dst) will be ignored.
-    
+
     Return:
         A list of subgraphs
     """
@@ -476,3 +476,43 @@ def pinsage_sample(graph,
             layer_nodes[0], dtype="int64")
 
     return subgraphs
+
+
+def extract_edges_from_nodes(graph, sample_nodes):
+    eids = graph_kernel.extract_edges_from_nodes(
+        graph.adj_src_index._indptr, graph.adj_src_index._sorted_v,
+        graph.adj_src_index._sorted_eid, sample_nodes)
+    return eids
+
+
+def graph_saint_random_walk_sample(graph,
+                                   nodes,
+                                   max_depth,
+                                   alias_name=None,
+                                   events_name=None):
+    """Implement of graph saint random walk sample.
+
+    First, this function will get random walks path for given nodes and depth.
+    Then, it will create subgraph from all sampled nodes.
+
+    Reference Paper: https://arxiv.org/abs/1907.04931
+
+    Args:
+        graph: A pgl graph instance
+        nodes: Walk starting from nodes
+        max_depth: Max walking depth
+
+    Return:
+        a subgraph of sampled nodes.
+    """
+    graph.outdegree()
+    walks = deepwalk_sample(graph, nodes, max_depth, alias_name, events_name)
+    sample_nodes = []
+    for walk in walks:
+        sample_nodes.extend(walk)
+    sample_nodes = np.unique(sample_nodes)
+    eids = extract_edges_from_nodes(graph, sample_nodes)
+    subgraph = graph.subgraph(
+        nodes=sample_nodes, eid=eids, with_node_feat=True, with_edge_feat=True)
+    subgraph.node_feat["index"] = np.array(sample_nodes, dtype="int64")
+    return subgraph

pgl/tests/test_graph_saint_sample.py

+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""graph saint sample test
+"""
+from __future__ import division
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import unicode_literals
+import unittest
+import numpy as np
+
+import pgl
+import paddle.fluid as fluid
+from pgl.sample import graph_saint_random_walk_sample
+
+
+class GraphSaintSampleTest(unittest.TestCase):
+    """GraphSaintSampleTest"""
+
+    def test_randomwalk_sampler(self):
+        """test_randomwalk_sampler"""
+        g = pgl.graph.Graph(
+            num_nodes=8,
+            edges=[(1, 2), (2, 3), (0, 2), (0, 1), (6, 7), (4, 5), (6, 4),
+                   (7, 4), (3, 4)])
+        subgraph = graph_saint_random_walk_sample(g, [6, 7], 2)
+        print('reindex', subgraph._from_reindex)
+        print('subedges', subgraph.edges)
+        assert len(subgraph.nodes) == 4
+        assert len(subgraph.edges) == 4
+        true_edges = np.array([[0, 1], [2, 3], [2, 0], [3, 0]])
+        assert "{}".format(subgraph.edges) == "{}".format(true_edges)
+
+
+if __name__ == '__main__':
+    unittest.main()