[Sampling] NodeDataLoader for node classification (#1635)

* neighbor sampler data loader first commit

* more commit

* nodedataloader

* fix

* update RGCN example

* update OGB

* fixes

* fix minibatch RGCN crashing with self loop

* reverting gatconv test code

* fix

* change to new solution that doesn't require tf dataloader

* fix

* lint

* fix

* fixes

* change doc

* fix docstring

* docstring fixes

* return seeds and input nodes from data loader

* fixes

* fix test

* fix windows build problem

* add pytorch wrapper

* fixes

* add pytorch wrapper

* add unit test

* add -1 support to sample_neighbors & fix docstrings

* docstring fix

* lint

* add minibatch rgcn evaluations
Co-authored-by: xiang song(charlie.song) <classicxsong@gmail.com>
Co-authored-by: Tong He <hetong007@gmail.com>

python/dgl/sampling/dataloader.py

+"""Data loaders"""
+
+from collections.abc import Mapping
+from abc import ABC, abstractproperty, abstractmethod
+from .. import transform
+from ..base import NID, EID
+from .. import utils
+
+# pylint: disable=unused-argument
+def assign_block_eids(block, frontier, block_id, g, seed_nodes, *args, **kwargs):
+    """Assigns edge IDs from the original graph to the block.
+
+    This is the default block postprocessor for samplers created with
+    ``return_eids`` as True.
+
+    See also
+    --------
+    BlockSampler
+    MultiLayerNeighborSampler
+    """
+    for etype in block.canonical_etypes:
+        block.edges[etype].data[EID] = frontier.edges[etype].data[EID][
+            block.edges[etype].data[EID]]
+    return block
+
+def _default_frontier_postprocessor(frontier, block_id, g, seed_nodes, *args, **kwargs):
+    return frontier
+
+def _default_block_postprocessor(block, frontier, block_id, g, seed_nodes, *args, **kwargs):
+    return block
+
+class BlockSampler(object):
+    """Abstract class specifying the neighborhood sampling strategy for DGL data loaders.
+
+    The main method for BlockSampler is :func:`~dgl.sampling.BlockSampler.sample_blocks`,
+    which generates a list of blocks for a multi-layer GNN given a set of seed nodes to
+    have their outputs computed.
+
+    The default implementation of :py:meth:`~dgl.sampling.BlockSampler.sample_blocks` is
+    to repeat ``num_hops`` times the following:
+
+    * Obtain a frontier with the same nodes as the original graph but only the edges
+      involved in message passing on the last layer.
+      Customizable via :py:meth:`~dgl.sampling.BlockSampler.sample_frontier`.
+
+    * Optionally, post-process the obtained frontier (e.g. by removing edges connecting training
+      node pairs).  One can add such postprocessors via
+      :py:meth:`~dgl.sampling.BlockSampler.add_frontier_postprocessor`.
+
+    * Convert the frontier into a block.
+
+    * Optionally, post-process the block (e.g. by assigning edge IDs).  One can add such
+      postprocessors via
+      :py:meth:`~dgl.sampling.BlockSampler.add_block_postprocessor`.
+
+    * Prepend the block to the block list to be returned.
+
+    All subclasses should either
+
+    * Override :py:meth:`~dgl.sampling.BlockSampler.sample_blocks` method, or
+
+    * Override
+      :py:meth:`~dgl.sampling.BlockSampler.sample_frontier` method while specifying
+      the number of layers to sample in ``num_hops`` argument.
+
+    See also
+    --------
+    For the concept of frontiers and blocks, please refer to User Guide Section 6.
+    """
+    def __init__(self, num_hops):
+        self.num_hops = num_hops
+        self._frontier_postprocessor = _default_frontier_postprocessor
+        self._block_postprocessor = _default_block_postprocessor
+
+    @property
+    def frontier_postprocessor(self):
+        """Frontier postprocessor."""
+        return self._frontier_postprocessor
+
+    @property
+    def block_postprocessor(self):
+        """B;pcl postprocessor."""
+        return self._block_postprocessor
+
+    def set_frontier_postprocessor(self, postprocessor):
+        """Set a frontier postprocessor.
+
+        The postprocessor must have the following signature:
+
+        .. code::
+
+           postprocessor(frontier, block_id, g, seed_nodes, *args, **kwargs)
+
+        where
+
+        * ``frontier`` represents the frontier obtained by
+          :py:meth:`~dgl.sampling.BlockSampler.sample_frontier` method.
+
+        * ``block_id`` represents which GNN layer the block is currently generated for.
+
+        * ``g`` represents the original graph.
+
+        * ``seed_nodes`` represents the output nodes on the current layer.
+
+        * Other arguments are the same ones passed into
+          :py:meth:`~dgl.sampling.BlockSampler.sample_blocks` method.
+
+        Parameters
+        ----------
+        postprocessor : callable
+            The postprocessor.
+        """
+        self._frontier_postprocessor = postprocessor
+
+    def set_block_postprocessor(self, postprocessor):
+        """Set a block postprocessor.
+
+        The postprocessor must have the following signature:
+
+        .. code::
+
+           postprocessor(block, frontier, block_id, g, seed_nodes, *args, **kwargs)
+
+        where
+
+        * ``block`` represents the block converted from the frontier.
+
+        * ``frontier`` represents the frontier the block is generated from.
+
+        * ``block_id`` represents which GNN layer the block is currently generated for.
+
+        * ``g`` represents the original graph.
+
+        * ``seed_nodes`` represents the output nodes on the current layer.
+
+        * Other arguments are the same ones passed into
+          :py:meth:`~dgl.sampling.BlockSampler.sample_blocks` method.
+
+        Parameters
+        ----------
+        postprocessor : callable
+            The postprocessor.
+        """
+        self._block_postprocessor = postprocessor
+
+    def _postprocess_frontier(self, frontier, block_id, g, seed_nodes, *args, **kwargs):
+        """Post-processes the generated frontier."""
+        return self._frontier_postprocessor(
+            frontier, block_id, g, seed_nodes, *args, **kwargs)
+
+    def _postprocess_block(self, block, frontier, block_id, g, seed_nodes, *args, **kwargs):
+        """Post-processes the generated block."""
+        return self._block_postprocessor(
+            block, frontier, block_id, g, seed_nodes, *args, **kwargs)
+
+    def sample_frontier(self, block_id, g, seed_nodes, *args, **kwargs):
+        """
+        Generate the frontier given the output nodes.
+
+        Parameters
+        ----------
+        block_id : int
+            Represents which GNN layer the frontier is generated for.
+        g : DGLHeteroGraph
+            The original graph.
+        seed_nodes : Tensor or dict[ntype, Tensor]
+            The output nodes by node type.
+
+            If the graph only has one node type, one can just specify a single tensor
+            of node IDs.
+        args, kwargs :
+            Other arguments being passed by
+            :py:meth:`~dgl.sampling.BlockSampler.sample_blocks`.
+
+        Returns
+        -------
+        DGLHeteroGraph
+            The frontier generated for the current layer.
+
+        See also
+        --------
+        For the concept of frontiers and blocks, please refer to User Guide Section 6.
+        """
+        raise NotImplementedError
+
+    def sample_blocks(self, g, seed_nodes, *args, **kwargs):
+        """
+        Generate the a list of blocks given the output nodes.
+
+        Parameters
+        ----------
+        g : DGLHeteroGraph
+            The original graph.
+        seed_nodes : Tensor or dict[ntype, Tensor]
+            The output nodes by node type.
+
+            If the graph only has one node type, one can just specify a single tensor
+            of node IDs.
+        args, kwargs :
+            Other arguments being passed by
+            :py:meth:`~dgl.sampling.BlockSampler.sample_blocks`.
+
+        Returns
+        -------
+        list[DGLHeteroGraph]
+            The blocks generated for computing the multi-layer GNN output.
+
+        See also
+        --------
+        For the concept of frontiers and blocks, please refer to User Guide Section 6.
+        """
+        blocks = []
+        for block_id in reversed(range(self.num_hops)):
+            frontier = self.sample_frontier(block_id, g, seed_nodes, *args, **kwargs)
+            # Removing edges from the frontier for link prediction training falls
+            # into the category of frontier postprocessing
+            frontier = self._postprocess_frontier(
+                frontier, block_id, g, seed_nodes, *args, **kwargs)
+
+            block = transform.to_block(frontier, seed_nodes)
+            # Assigning edge IDs and/or node/edge features falls into the category of block
+            # postprocessing
+            block = self._postprocess_block(
+                block, frontier, block_id, g, seed_nodes, *args, **kwargs)
+
+            seed_nodes = {ntype: block.srcnodes[ntype].data[NID] for ntype in block.srctypes}
+            blocks.insert(0, block)
+        return blocks
+
+class Collator(ABC):
+    """
+    Abstract DGL collator for training GNNs on downstream tasks stochastically.
+
+    Provides a ``dataset`` object containing the collection of all nodes or edges,
+    as well as a ``collate`` method that combines a set of items from ``dataset`` and
+    obtains the blocks.
+
+    See also
+    --------
+    For the concept of blocks, please refer to User Guide Section 6.
+    """
+    @abstractproperty
+    def dataset(self):
+        """Returns the dataset object of the collator."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def collate(self, items):
+        """Combines the items from the dataset object and obtains the list of blocks.
+
+        Parameters
+        ----------
+        items : list[str, int]
+            The list of node or edge type-ID pairs.
+
+        See also
+        --------
+        For the concept of blocks, please refer to User Guide Section 6.
+        """
+        raise NotImplementedError
+
+class NodeCollator(Collator):
+    """
+    DGL collator to combine training node classification or regression on a single graph.
+
+    Parameters
+    ----------
+    g : DGLHeteroGraph
+        The graph.
+    nids : Tensor or dict[ntype, Tensor]
+        The node set to compute outputs.
+    block_sampler : :py:class:`~dgl.sampling.BlockSampler`
+        The neighborhood sampler.
+
+    Examples
+    --------
+    To train a 3-layer GNN for node classification on a set of nodes ``train_nid`` on
+    a homogeneous graph where each node takes messages from all neighbors (assume
+    the backend is PyTorch):
+    >>> sampler = dgl.sampling.NeighborSampler([None, None, None])
+    >>> collator = dgl.sampling.NodeCollator(g, train_nid, sampler)
+    >>> dataloader = torch.utils.data.DataLoader(
+    ...     collator.dataset, collate_fn=collator.collate,
+    ...     batch_size=1024, shuffle=True, drop_last=False, num_workers=4)
+    >>> for input_nodes, output_nodes, blocks in dataloader:
+    ...     train_on(input_nodes, output_nodes, blocks)
+    """
+    def __init__(self, g, nids, block_sampler):
+        self.g = g
+        if not isinstance(nids, Mapping):
+            assert len(g.ntypes) == 1, \
+                "nids should be a dict of node type and ids for graph with multiple node types"
+        self.nids = nids
+        self.block_sampler = block_sampler
+
+        if isinstance(nids, Mapping):
+            self._dataset = utils.FlattenedDict(nids)
+        else:
+            self._dataset = nids
+
+    @property
+    def dataset(self):
+        return self._dataset
+
+    def collate(self, items):
+        """Find the list of blocks necessary for computing the representation of given
+        nodes for a node classification/regression task.
+
+        Returns
+        -------
+        input_nodes : Tensor or dict[ntype, Tensor]
+            The input nodes necessary for computation in this minibatch.
+
+            If the original graph has multiple node types, return a dictionary of
+            node type names and node ID tensors.  Otherwise, return a single tensor.
+        output_nodes : Tensor or dict[ntype, Tensor]
+            The nodes whose representations are to be computed in this minibatch.
+
+            If the original graph has multiple node types, return a dictionary of
+            node type names and node ID tensors.  Otherwise, return a single tensor.
+        blocks : list[DGLHeteroGraph]
+            The list of blocks necessary for computing the representation.
+        """
+        if isinstance(items[0], tuple):
+            # returns a list of pairs: group them by node types into a dict
+            items = utils.group_as_dict(items)
+        blocks = self.block_sampler.sample_blocks(self.g, items)
+
+        if len(self.g.ntypes) == 1:
+            output_nodes = blocks[-1].dstdata[NID]
+            input_nodes = blocks[0].srcdata[NID]
+        else:
+            output_nodes = {
+                ntype: blocks[-1].dstnodes[ntype].data[NID]
+                for ntype in blocks[-1].dsttypes}
+            input_nodes = {
+                ntype: blocks[0].srcnodes[ntype].data[NID]
+                for ntype in blocks[0].srctypes}
+
+        return input_nodes, output_nodes, blocks

python/dgl/sampling/neighbor.py

@@ -6,10 +6,13 @@ from ..base import DGLError, EID
 from ..heterograph import DGLHeteroGraph
 from .. import ndarray as nd
 from .. import utils
+from .. import transform
+from .dataloader import BlockSampler, assign_block_eids
 
 __all__ = [
     'sample_neighbors',
-    'select_topk']
+    'select_topk',
+    'MultiLayerNeighborSampler']
 
 def sample_neighbors(g, nodes, fanout, edge_dir='in', prob=None, replace=False):
     """Sample from the neighbors of the given nodes and return the induced subgraph.
@@ -33,6 +36,9 @@ def sample_neighbors(g, nodes, fanout, edge_dir='in', prob=None, replace=False):
     fanout : int or dict[etype, int]
         The number of sampled neighbors for each node on each edge type. Provide a dict
         to specify different fanout values for each edge type.
+
+        If -1 is given, select all the neighbors.  ``prob`` and ``replace`` will be
+        ignored in this case.
     edge_dir : str, optional
         Edge direction ('in' or 'out'). If is 'in', sample from in edges. Otherwise,
         sample from out edges.
@@ -56,7 +62,7 @@ def sample_neighbors(g, nodes, fanout, edge_dir='in', prob=None, replace=False):
     nodes_all_types = []
     for ntype in g.ntypes:
         if ntype in nodes:
-            nodes_all_types.append(utils.toindex(nodes[ntype]).todgltensor())
+            nodes_all_types.append(utils.toindex(nodes[ntype], g._idtype_str).todgltensor())
         else:
             nodes_all_types.append(nd.array([], ctx=nd.cpu()))
 
@@ -103,6 +109,8 @@ def select_topk(g, k, weight, nodes=None, edge_dir='in', ascending=False):
         Full graph structure.
     k : int or dict[etype, int]
         The K value.
+
+        If -1 is given, select all the neighbors.
     weight : str
         Feature name of the weights associated with each edge. Its shape should be
         compatible with a scalar edge feature tensor.
@@ -135,7 +143,7 @@ def select_topk(g, k, weight, nodes=None, edge_dir='in', ascending=False):
     nodes_all_types = []
     for ntype in g.ntypes:
         if ntype in nodes:
-            nodes_all_types.append(utils.toindex(nodes[ntype]).todgltensor())
+            nodes_all_types.append(utils.toindex(nodes[ntype], g._idtype_str).todgltensor())
         else:
             nodes_all_types.append(nd.array([], ctx=nd.cpu()))
 
@@ -166,4 +174,74 @@ def select_topk(g, k, weight, nodes=None, edge_dir='in', ascending=False):
         ret.edges[etype].data[EID] = induced_edges[i].tousertensor()
     return ret
 
+
+class MultiLayerNeighborSampler(BlockSampler):
+    """Sampler that builds computational dependency of node representations via
+    neighbor sampling for multilayer GNN.
+
+    This sampler will make every node gather messages from a fixed number of neighbors
+    per edge type.  The neighbors are picked uniformly.
+
+    Parameters
+    ----------
+    fanouts : list[int] or list[dict[etype, int] or None]
+        List of neighbors to sample per edge type for each GNN layer, starting from the
+        first layer.
+
+        If the graph is homogeneous, only an integer is needed for each layer.
+
+        If None is provided for one layer, all neighbors will be included regardless of
+        edge types.
+
+        If -1 is provided for one edge type on one layer, then all inbound edges
+        of that edge type will be included.
+    replace : bool, default True
+        Whether to sample with replacement
+    return_eids : bool, default False
+        Whether to return edge IDs of the original graph in the sampled blocks.
+
+        If True, the edge IDs will be stored as ``dgl.EID`` feature for each edge type.
+
+    Examples
+    --------
+    To train a 3-layer GNN for node classification on a set of nodes ``train_nid`` on
+    a homogeneous graph where each node takes messages from all neighbors (assume
+    the backend is PyTorch):
+    >>> sampler = dgl.sampling.NeighborSampler([None, None, None])
+    >>> collator = dgl.sampling.NodeCollator(g, train_nid, sampler)
+    >>> dataloader = torch.utils.data.DataLoader(
+    ...     collator.dataset, collate_fn=collator.collate,
+    ...     batch_size=1024, shuffle=True, drop_last=False, num_workers=4)
+    >>> for blocks in dataloader:
+    ...     train_on(blocks)
+
+    If we wish to gather from 5 neighbors on the first layer, 10 neighbors on the second,
+    and 15 layers on the third:
+    >>> sampler = dgl.sampling.NeighborSampler([5, 10, 15])
+
+    If training on a heterogeneous graph and you want different number of neighbors for each
+    edge type, one should instead provide a list of dicts.  Each dict would specify the
+    number of neighbors to pick per edge type.
+    >>> sampler = dgl.sampling.NeighborSampler([
+    ...     {('user', 'follows', 'user'): 5,
+    ...      ('user', 'plays', 'game'): 4,
+    ...      ('game', 'played-by', 'user'): 3}] * 3)
+    """
+    def __init__(self, fanouts, replace=False, return_eids=False):
+        super().__init__(len(fanouts))
+
+        self.fanouts = fanouts
+        self.replace = replace
+        self.return_eids = return_eids
+        if return_eids:
+            self.set_block_postprocessor(assign_block_eids)
+
+    def sample_frontier(self, block_id, g, seed_nodes, *args, **kwargs):
+        fanout = self.fanouts[block_id]
+        if fanout is None:
+            frontier = transform.in_subgraph(g, seed_nodes)
+        else:
+            frontier = sample_neighbors(g, seed_nodes, fanout, replace=self.replace)
+        return frontier
+
 _init_api('dgl.sampling.neighbor', __name__)

python/dgl/sampling/pytorch/__init__.py

+"""DGL PyTorch DataLoaders"""
+from torch.utils.data import DataLoader
+from ..dataloader import NodeCollator
+
+class NodeDataLoader(DataLoader):
+    """PyTorch dataloader for batch-iterating over a set of nodes, generating the list
+    of blocks as computation dependency of the said minibatch.
+
+    Parameters
+    ----------
+    g : DGLHeteroGraph
+        The graph.
+    nids : Tensor or dict[ntype, Tensor]
+        The node set to compute outputs.
+    block_sampler : :py:class:`~dgl.sampling.BlockSampler`
+        The neighborhood sampler.
+    kwargs : dict
+        Arguments being passed to `torch.utils.data.DataLoader`.
+
+    Examples
+    --------
+    To train a 3-layer GNN for node classification on a set of nodes ``train_nid`` on
+    a homogeneous graph where each node takes messages from all neighbors (assume
+    the backend is PyTorch):
+    >>> sampler = dgl.sampling.NeighborSampler([None, None, None])
+    >>> dataloader = dgl.sampling.NodeDataLoader(
+    ...     g, train_nid, sampler,
+    ...     batch_size=1024, shuffle=True, drop_last=False, num_workers=4)
+    >>> for input_nodes, output_nodes, blocks in dataloader:
+    ...     train_on(input_nodes, output_nodes, blocks)
+    """
+    def __init__(self, g, nids, block_sampler, **kwargs):
+        self.collator = NodeCollator(g, nids, block_sampler)
+        super().__init__(self.collator.dataset, collate_fn=self.collator.collate, **kwargs)

python/dgl/sampling/randomwalks.py

@@ -123,8 +123,8 @@ def random_walk(g, nodes, *, metapath=None, length=None, prob=None, restart_prob
         metapath = [g.get_etype_id(etype) for etype in metapath]
 
     gidx = g._graph
-    nodes = utils.toindex(nodes).todgltensor()
-    metapath = utils.toindex(metapath).todgltensor().copyto(nodes.ctx)
+    nodes = utils.toindex(nodes, g._idtype_str).todgltensor()
+    metapath = utils.toindex(metapath, g._idtype_str).todgltensor().copyto(nodes.ctx)
 
     # Load the probability tensor from the edge frames
     if prob is None:

python/dgl/transform.py

@@ -1044,6 +1044,9 @@ def to_block(g, dst_nodes=None, include_dst_in_src=True):
             raise ValueError(
                 'Graph has more than one node type; please specify a dict for dst_nodes.')
         dst_nodes = {g.ntypes[0]: dst_nodes}
+    dst_nodes = {
+        ntype: utils.toindex(nodes, g._idtype_str).tousertensor()
+        for ntype, nodes in dst_nodes.items()}
 
     # dst_nodes is now a dict
     dst_nodes_nd = []

python/dgl/utils.py

@@ -2,6 +2,7 @@
 from __future__ import absolute_import, division
 
 from collections.abc import Mapping, Iterable
+from collections import defaultdict
 from functools import wraps
 import numpy as np
 
@@ -598,3 +599,64 @@ def retry_method_with_fix(fix_method):
 
         return wrapper
     return _creator
+
+def group_as_dict(pairs):
+    """Combines a list of key-value pairs to a dictionary of keys and value lists.
+
+    Does not require the pairs to be sorted by keys.
+
+    Parameters
+    ----------
+    pairs : iterable
+        Iterable of key-value pairs
+
+    Returns
+    -------
+    dict
+        The dictionary of keys and value lists.
+    """
+    dic = defaultdict(list)
+    for key, value in pairs:
+        dic[key].append(value)
+    return dic
+
+class FlattenedDict(object):
+    """Iterates over each item in a dictionary of groups.
+
+    Parameters
+    ----------
+    groups : dict
+        The item groups.
+
+    Examples
+    --------
+    >>> groups = FlattenedDict({'a': [1, 3], 'b': [2, 5, 8], 'c': [7]})
+    >>> list(groups)
+    [('a', 1), ('a', 3), ('b', 2), ('b', 5), ('b', 8), ('c', 7)]
+    >>> groups[2]
+    ('b', 2)
+    >>> len(groups)
+    6
+    """
+    def __init__(self, groups):
+        self._groups = groups
+        group_sizes = {k: len(v) for k, v in groups.items()}
+        self._group_keys, self._group_sizes = zip(*group_sizes.items())
+        self._group_offsets = np.insert(np.cumsum(self._group_sizes), 0, 0)
+
+    def __len__(self):
+        """Return the total number of items."""
+        return self._group_offsets[-1]
+
+    def __iter__(self):
+        """Return the iterator of all items with the key of its original group."""
+        for i, k in enumerate(self._group_keys):
+            for j in range(self._group_sizes[i]):
+                yield k, self._groups[k][j]
+
+    def __getitem__(self, idx):
+        """Return the item at the given position with the key of its original group."""
+        i = np.searchsorted(self._group_offsets, idx, 'right') - 1
+        k = self._group_keys[i]
+        j = idx - self._group_offsets[i]
+        return k, self._groups[k][j]

src/graph/sampling/neighbor/neighbor.cc

@@ -49,6 +49,17 @@ HeteroSubgraph SampleNeighbors(
         hg->NumVertices(dst_vtype),
         hg->DataType(), hg->Context());
       induced_edges[etype] = aten::NullArray();
+    } else if (fanouts[etype] == -1) {
+      const auto &earr = (dir == EdgeDir::kOut) ?
+        hg->OutEdges(etype, nodes_ntype) :
+        hg->InEdges(etype, nodes_ntype);
+      subrels[etype] = UnitGraph::CreateFromCOO(
+        hg->GetRelationGraph(etype)->NumVertexTypes(),
+        hg->NumVertices(src_vtype),
+        hg->NumVertices(dst_vtype),
+        earr.src,
+        earr.dst);
+      induced_edges[etype] = earr.id;
     } else {
       // sample from one relation graph
       auto req_fmt = (dir == EdgeDir::kOut)? SparseFormat::kCSR : SparseFormat::kCSC;
@@ -124,6 +135,17 @@ HeteroSubgraph SampleNeighborsTopk(
         hg->NumVertices(dst_vtype),
         hg->DataType(), hg->Context());
       induced_edges[etype] = aten::NullArray();
+    } else if (k[etype] == -1) {
+      const auto &earr = (dir == EdgeDir::kOut) ?
+        hg->OutEdges(etype, nodes_ntype) :
+        hg->InEdges(etype, nodes_ntype);
+      subrels[etype] = UnitGraph::CreateFromCOO(
+        hg->GetRelationGraph(etype)->NumVertexTypes(),
+        hg->NumVertices(src_vtype),
+        hg->NumVertices(dst_vtype),
+        earr.src,
+        earr.dst);
+      induced_edges[etype] = earr.id;
     } else {
       // sample from one relation graph
       auto req_fmt = (dir == EdgeDir::kOut)? SparseFormat::kCSR : SparseFormat::kCSC;

tests/compute/test_sampling.py

@@ -2,6 +2,7 @@ import dgl
 import backend as F
 import numpy as np
 import unittest
+from collections import defaultdict
 
 def check_random_walk(g, metapath, traces, ntypes, prob=None):
     traces = F.asnumpy(traces)
@@ -213,6 +214,14 @@ def _test_sample_neighbors(hypersparse):
     g, hg = _gen_neighbor_sampling_test_graph(hypersparse, False)
 
     def _test1(p, replace):
+        subg = dgl.sampling.sample_neighbors(g, [0, 1], -1, prob=p, replace=replace)
+        assert subg.number_of_nodes() == g.number_of_nodes()
+        u, v = subg.edges()
+        u_ans, v_ans = subg.in_edges([0, 1])
+        uv = set(zip(F.asnumpy(u), F.asnumpy(v)))
+        uv_ans = set(zip(F.asnumpy(u_ans), F.asnumpy(v_ans)))
+        assert uv == uv_ans
+
         for i in range(10):
             subg = dgl.sampling.sample_neighbors(g, [0, 1], 2, prob=p, replace=replace)
             assert subg.number_of_nodes() == g.number_of_nodes()
@@ -234,6 +243,14 @@ def _test_sample_neighbors(hypersparse):
     _test1('prob', False)  # w/o replacement
 
     def _test2(p, replace):  # fanout > #neighbors
+        subg = dgl.sampling.sample_neighbors(g, [0, 2], -1, prob=p, replace=replace)
+        assert subg.number_of_nodes() == g.number_of_nodes()
+        u, v = subg.edges()
+        u_ans, v_ans = subg.in_edges([0, 2])
+        uv = set(zip(F.asnumpy(u), F.asnumpy(v)))
+        uv_ans = set(zip(F.asnumpy(u_ans), F.asnumpy(v_ans)))
+        assert uv == uv_ans
+
         for i in range(10):
             subg = dgl.sampling.sample_neighbors(g, [0, 2], 2, prob=p, replace=replace)
             assert subg.number_of_nodes() == g.number_of_nodes()
@@ -255,6 +272,14 @@ def _test_sample_neighbors(hypersparse):
     _test2('prob', False)  # w/o replacement
 
     def _test3(p, replace):
+        subg = dgl.sampling.sample_neighbors(hg, {'user': [0, 1], 'game': 0}, -1, prob=p, replace=replace)
+        assert len(subg.ntypes) == 3
+        assert len(subg.etypes) == 4
+        assert subg['follow'].number_of_edges() == 6
+        assert subg['play'].number_of_edges() == 1
+        assert subg['liked-by'].number_of_edges() == 4
+        assert subg['flips'].number_of_edges() == 0
+
         for i in range(10):
             subg = dgl.sampling.sample_neighbors(hg, {'user' : [0,1], 'game' : 0}, 2, prob=p, replace=replace)
             assert len(subg.ntypes) == 3
@@ -273,20 +298,28 @@ def _test_sample_neighbors(hypersparse):
     for i in range(10):
         subg = dgl.sampling.sample_neighbors(
             hg,
-            {'user' : [0,1], 'game' : 0},
-            {'follow': 1, 'play': 2, 'liked-by': 0, 'flips': 2},
+            {'user' : [0,1], 'game' : 0, 'coin': 0},
+            {'follow': 1, 'play': 2, 'liked-by': 0, 'flips': -1},
             replace=True)
         assert len(subg.ntypes) == 3
         assert len(subg.etypes) == 4
         assert subg['follow'].number_of_edges() == 2
         assert subg['play'].number_of_edges() == 2
         assert subg['liked-by'].number_of_edges() == 0
-        assert subg['flips'].number_of_edges() == 0
+        assert subg['flips'].number_of_edges() == 4
 
 def _test_sample_neighbors_outedge(hypersparse):
     g, hg = _gen_neighbor_sampling_test_graph(hypersparse, True)
 
     def _test1(p, replace):
+        subg = dgl.sampling.sample_neighbors(g, [0, 1], -1, prob=p, replace=replace, edge_dir='out')
+        assert subg.number_of_nodes() == g.number_of_nodes()
+        u, v = subg.edges()
+        u_ans, v_ans = subg.out_edges([0, 1])
+        uv = set(zip(F.asnumpy(u), F.asnumpy(v)))
+        uv_ans = set(zip(F.asnumpy(u_ans), F.asnumpy(v_ans)))
+        assert uv == uv_ans
+
         for i in range(10):
             subg = dgl.sampling.sample_neighbors(g, [0, 1], 2, prob=p, replace=replace, edge_dir='out')
             assert subg.number_of_nodes() == g.number_of_nodes()
@@ -308,6 +341,14 @@ def _test_sample_neighbors_outedge(hypersparse):
     _test1('prob', False)  # w/o replacement
 
     def _test2(p, replace):  # fanout > #neighbors
+        subg = dgl.sampling.sample_neighbors(g, [0, 2], -1, prob=p, replace=replace, edge_dir='out')
+        assert subg.number_of_nodes() == g.number_of_nodes()
+        u, v = subg.edges()
+        u_ans, v_ans = subg.out_edges([0, 2])
+        uv = set(zip(F.asnumpy(u), F.asnumpy(v)))
+        uv_ans = set(zip(F.asnumpy(u_ans), F.asnumpy(v_ans)))
+        assert uv == uv_ans
+
         for i in range(10):
             subg = dgl.sampling.sample_neighbors(g, [0, 2], 2, prob=p, replace=replace, edge_dir='out')
             assert subg.number_of_nodes() == g.number_of_nodes()
@@ -329,6 +370,14 @@ def _test_sample_neighbors_outedge(hypersparse):
     _test2('prob', False)  # w/o replacement
 
     def _test3(p, replace):
+        subg = dgl.sampling.sample_neighbors(hg, {'user': [0, 1], 'game': 0}, -1, prob=p, replace=replace, edge_dir='out')
+        assert len(subg.ntypes) == 3
+        assert len(subg.etypes) == 4
+        assert subg['follow'].number_of_edges() == 6
+        assert subg['play'].number_of_edges() == 1
+        assert subg['liked-by'].number_of_edges() == 4
+        assert subg['flips'].number_of_edges() == 0
+
         for i in range(10):
             subg = dgl.sampling.sample_neighbors(hg, {'user' : [0,1], 'game' : 0}, 2, prob=p, replace=replace, edge_dir='out')
             assert len(subg.ntypes) == 3
@@ -347,6 +396,14 @@ def _test_sample_neighbors_topk(hypersparse):
     g, hg = _gen_neighbor_topk_test_graph(hypersparse, False)
 
     def _test1():
+        subg = dgl.sampling.select_topk(g, -1, 'weight', [0, 1])
+        assert subg.number_of_nodes() == g.number_of_nodes()
+        u, v = subg.edges()
+        u_ans, v_ans = subg.in_edges([0, 1])
+        uv = set(zip(F.asnumpy(u), F.asnumpy(v)))
+        uv_ans = set(zip(F.asnumpy(u_ans), F.asnumpy(v_ans)))
+        assert uv == uv_ans
+
         subg = dgl.sampling.select_topk(g, 2, 'weight', [0, 1])
         assert subg.number_of_nodes() == g.number_of_nodes()
         assert subg.number_of_edges() == 4
@@ -357,6 +414,14 @@ def _test_sample_neighbors_topk(hypersparse):
     _test1()
 
     def _test2():  # k > #neighbors
+        subg = dgl.sampling.select_topk(g, -1, 'weight', [0, 2])
+        assert subg.number_of_nodes() == g.number_of_nodes()
+        u, v = subg.edges()
+        u_ans, v_ans = subg.in_edges([0, 2])
+        uv = set(zip(F.asnumpy(u), F.asnumpy(v)))
+        uv_ans = set(zip(F.asnumpy(u_ans), F.asnumpy(v_ans)))
+        assert uv == uv_ans
+
         subg = dgl.sampling.select_topk(g, 2, 'weight', [0, 2])
         assert subg.number_of_nodes() == g.number_of_nodes()
         assert subg.number_of_edges() == 3
@@ -387,18 +452,26 @@ def _test_sample_neighbors_topk(hypersparse):
 
     # test different k for different relations
     subg = dgl.sampling.select_topk(
-        hg, {'follow': 1, 'play': 2, 'liked-by': 0, 'flips': 2}, 'weight', {'user' : [0,1], 'game' : 0})
+        hg, {'follow': 1, 'play': 2, 'liked-by': 0, 'flips': -1}, 'weight', {'user' : [0,1], 'game' : 0, 'coin': 0})
     assert len(subg.ntypes) == 3
     assert len(subg.etypes) == 4
     assert subg['follow'].number_of_edges() == 2
     assert subg['play'].number_of_edges() == 1
     assert subg['liked-by'].number_of_edges() == 0
-    assert subg['flips'].number_of_edges() == 0
+    assert subg['flips'].number_of_edges() == 4
 
 def _test_sample_neighbors_topk_outedge(hypersparse):
     g, hg = _gen_neighbor_topk_test_graph(hypersparse, True)
 
     def _test1():
+        subg = dgl.sampling.select_topk(g, -1, 'weight', [0, 1], edge_dir='out')
+        assert subg.number_of_nodes() == g.number_of_nodes()
+        u, v = subg.edges()
+        u_ans, v_ans = subg.out_edges([0, 1])
+        uv = set(zip(F.asnumpy(u), F.asnumpy(v)))
+        uv_ans = set(zip(F.asnumpy(u_ans), F.asnumpy(v_ans)))
+        assert uv == uv_ans
+
         subg = dgl.sampling.select_topk(g, 2, 'weight', [0, 1], edge_dir='out')
         assert subg.number_of_nodes() == g.number_of_nodes()
         assert subg.number_of_edges() == 4
@@ -409,6 +482,14 @@ def _test_sample_neighbors_topk_outedge(hypersparse):
     _test1()
 
     def _test2():  # k > #neighbors
+        subg = dgl.sampling.select_topk(g, -1, 'weight', [0, 2], edge_dir='out')
+        assert subg.number_of_nodes() == g.number_of_nodes()
+        u, v = subg.edges()
+        u_ans, v_ans = subg.out_edges([0, 2])
+        uv = set(zip(F.asnumpy(u), F.asnumpy(v)))
+        uv_ans = set(zip(F.asnumpy(u_ans), F.asnumpy(v_ans)))
+        assert uv == uv_ans
+
         subg = dgl.sampling.select_topk(g, 2, 'weight', [0, 2], edge_dir='out')
         assert subg.number_of_nodes() == g.number_of_nodes()
         assert subg.number_of_edges() == 3

tests/mxnet/test_nn.py

@@ -86,11 +86,21 @@ def test_graph_conv2(g, norm, weight, bias):
     nsrc = g.number_of_nodes() if isinstance(g, dgl.DGLGraph) else g.number_of_src_nodes()
     ndst = g.number_of_nodes() if isinstance(g, dgl.DGLGraph) else g.number_of_dst_nodes()
     h = F.randn((nsrc, 5)).as_in_context(F.ctx())
+    h_dst = F.randn((ndst, 2)).as_in_context(F.ctx())
     if weight:
-        h = conv(g, h)
+        h_out = conv(g, h)
     else:
-        h = conv(g, h, ext_w)
-    assert h.shape == (ndst, 2)
+        h_out = conv(g, h, ext_w)
+    assert h_out.shape == (ndst, 2)
+
+    if not isinstance(g, dgl.DGLGraph) and len(g.ntypes) == 2:
+        # bipartite, should also accept pair of tensors
+        if weight:
+            h_out2 = conv(g, (h, h_dst))
+        else:
+            h_out2 = conv(g, (h, h_dst), ext_w)
+        assert h_out2.shape == (ndst, 2)
+        assert F.array_equal(h_out, h_out2)
 
 def _S2AXWb(A, N, X, W, b):
     X1 = X * N

tests/pytorch/test_dataloader.py

+import dgl
+import backend as F
+import numpy as np
+import unittest
+from torch.utils.data import DataLoader
+from collections import defaultdict
+
+def _check_neighbor_sampling_dataloader(g, nids, dl):
+    seeds = defaultdict(list)
+
+    for input_nodes, output_nodes, blocks in dl:
+        if len(g.ntypes) > 1:
+            for ntype in g.ntypes:
+                assert F.array_equal(input_nodes[ntype], blocks[0].srcnodes[ntype].data[dgl.NID])
+                assert F.array_equal(output_nodes[ntype], blocks[-1].dstnodes[ntype].data[dgl.NID])
+        else:
+            assert F.array_equal(input_nodes, blocks[0].srcdata[dgl.NID])
+            assert F.array_equal(output_nodes, blocks[-1].dstdata[dgl.NID])
+        prev_dst = {ntype: None for ntype in g.ntypes}
+        for block in blocks:
+            for canonical_etype in block.canonical_etypes:
+                utype, etype, vtype = canonical_etype
+                uu, vv = block.all_edges(order='eid', etype=canonical_etype)
+                src = block.srcnodes[utype].data[dgl.NID]
+                dst = block.dstnodes[vtype].data[dgl.NID]
+                if prev_dst[utype] is not None:
+                    assert F.array_equal(src, prev_dst[utype])
+                u = src[uu]
+                v = dst[vv]
+                assert F.asnumpy(g.has_edges_between(u, v, etype=canonical_etype)).all()
+                eid = block.edges[canonical_etype].data[dgl.EID]
+                ufound, vfound = g.find_edges(eid, etype=canonical_etype)
+                assert F.array_equal(ufound, u)
+                assert F.array_equal(vfound, v)
+            for ntype in block.dsttypes:
+                src = block.srcnodes[ntype].data[dgl.NID]
+                dst = block.dstnodes[ntype].data[dgl.NID]
+                assert F.array_equal(src[:block.number_of_dst_nodes(ntype)], dst)
+                prev_dst[ntype] = dst
+        for ntype in blocks[-1].dsttypes:
+            seeds[ntype].append(blocks[-1].dstnodes[ntype].data[dgl.NID])
+
+    # Check if all nodes are iterated
+    seeds = {k: F.cat(v, 0) for k, v in seeds.items()}
+    for k, v in seeds.items():
+        v_set = set(F.asnumpy(v))
+        seed_set = set(nids[k])
+        assert v_set == seed_set
+
+@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU sample neighbors not implemented")
+def test_neighbor_sampler_dataloader():
+    g = dgl.graph([(0,1),(0,2),(0,3),(1,0),(1,2),(1,3),(2,0)],
+            'user', 'follow', num_nodes=6)
+    g_sampler1 = dgl.sampling.MultiLayerNeighborSampler([2, 2], return_eids=True)
+    g_sampler2 = dgl.sampling.MultiLayerNeighborSampler([None, None], return_eids=True)
+
+    hg = dgl.heterograph({
+        ('user', 'follow', 'user'): [(0, 1), (0, 2), (0, 3), (1, 0), (1, 2), (1, 3), (2, 0)],
+        ('user', 'plays', 'game'): [(0, 0), (1, 1), (1, 2), (3, 0), (5, 2)],
+        ('game', 'wanted-by', 'user'): [(0, 1), (2, 1), (1, 3), (2, 3), (2, 5)]})
+    hg_sampler1 = dgl.sampling.MultiLayerNeighborSampler(
+        [{'plays': 1, 'wanted-by': 1, 'follow': 2}] * 2,
+        return_eids=True)
+    hg_sampler2 = dgl.sampling.MultiLayerNeighborSampler([None, None], return_eids=True)
+
+    collators = [
+        dgl.sampling.NodeCollator(g, [0, 1, 2, 3, 5], g_sampler1),
+        dgl.sampling.NodeCollator(g, [4, 5], g_sampler1),
+        dgl.sampling.NodeCollator(g, [0, 1, 2, 3, 5], g_sampler2),
+        dgl.sampling.NodeCollator(g, [4, 5], g_sampler2),
+        dgl.sampling.NodeCollator(hg, {'user': [0, 1, 3, 5], 'game': [0, 1, 2]}, hg_sampler1),
+        dgl.sampling.NodeCollator(hg, {'user': [4, 5], 'game': [0, 1, 2]}, hg_sampler1),
+        dgl.sampling.NodeCollator(hg, {'user': [0, 1, 3, 5], 'game': [0, 1, 2]}, hg_sampler2),
+        dgl.sampling.NodeCollator(hg, {'user': [4, 5], 'game': [0, 1, 2]}, hg_sampler2)]
+    nids = [
+        {'user': [0, 1, 2, 3, 5]},
+        {'user': [4, 5]},
+        {'user': [0, 1, 2, 3, 5]},
+        {'user': [4, 5]},
+        {'user': [0, 1, 3, 5], 'game': [0, 1, 2]},
+        {'user': [4, 5], 'game': [0, 1, 2]},
+        {'user': [0, 1, 3, 5], 'game': [0, 1, 2]},
+        {'user': [4, 5], 'game': [0, 1, 2]}]
+    graphs = [g] * 4 + [hg] * 4
+    samplers = [g_sampler1, g_sampler1, g_sampler2, g_sampler2, hg_sampler1, hg_sampler1, hg_sampler2, hg_sampler2]
+
+    for _g, nid, collator in zip(graphs, nids, collators):
+        dl = DataLoader(
+            collator.dataset, collate_fn=collator.collate, batch_size=2, shuffle=True, drop_last=False)
+        _check_neighbor_sampling_dataloader(_g, nid, dl)
+    for _g, nid, sampler in zip(graphs, nids, samplers):
+        dl = dgl.sampling.NodeDataLoader(_g, nid, sampler, batch_size=2, shuffle=True, drop_last=False)
+        _check_neighbor_sampling_dataloader(_g, nid, dl)
+
+
+if __name__ == '__main__':
+    test_neighbor_sampler_dataloader()

tests/pytorch/test_nn.py

@@ -79,11 +79,21 @@ def test_graph_conv2(g, norm, weight, bias):
     nsrc = g.number_of_nodes() if isinstance(g, dgl.DGLGraph) else g.number_of_src_nodes()
     ndst = g.number_of_nodes() if isinstance(g, dgl.DGLGraph) else g.number_of_dst_nodes()
     h = F.randn((nsrc, 5)).to(F.ctx())
+    h_dst = F.randn((ndst, 2)).to(F.ctx())
     if weight:
-        h = conv(g, h)
+        h_out = conv(g, h)
     else:
-        h = conv(g, h, weight=ext_w)
-    assert h.shape == (ndst, 2)
+        h_out = conv(g, h, weight=ext_w)
+    assert h_out.shape == (ndst, 2)
+
+    if not isinstance(g, dgl.DGLGraph) and len(g.ntypes) == 2:
+        # bipartite, should also accept pair of tensors
+        if weight:
+            h_out2 = conv(g, (h, h_dst))
+        else:
+            h_out2 = conv(g, (h, h_dst), weight=ext_w)
+        assert h_out2.shape == (ndst, 2)
+        assert F.array_equal(h_out, h_out2)
 
 def _S2AXWb(A, N, X, W, b):
     X1 = X * N

tests/scripts/build_dgl.bat

@@ -14,7 +14,7 @@ SET TMPDIR=%WORKSPACE%\\tmp
 PUSHD build
 CALL "C:\Program Files (x86)\Microsoft Visual Studio\2017\BuildTools\VC\Auxiliary\Build\vcvars64.bat"
 cmake -DCMAKE_CXX_FLAGS="/DDGL_EXPORTS" -DUSE_OPENMP=ON -Dgtest_force_shared_crt=ON -DDMLC_FORCE_SHARED_CRT=ON -DBUILD_CPP_TEST=1 -DCMAKE_CONFIGURATION_TYPES="Release" .. -G "Visual Studio 15 2017 Win64" || EXIT /B 1
-msbuild dgl.sln /m || EXIT /B 1
+msbuild dgl.sln /m /nr:false || EXIT /B 1
 COPY Release\dgl.dll .
 COPY Release\runUnitTests.exe .
 POPD

tests/tensorflow/test_nn.py

@@ -80,11 +80,21 @@ def test_graph_conv2(g, norm, weight, bias):
     nsrc = g.number_of_nodes() if isinstance(g, dgl.DGLGraph) else g.number_of_src_nodes()
     ndst = g.number_of_nodes() if isinstance(g, dgl.DGLGraph) else g.number_of_dst_nodes()
     h = F.randn((nsrc, 5))
+    h_dst = F.randn((ndst, 2))
     if weight:
-        h = conv(g, h)
+        h_out = conv(g, h)
     else:
-        h = conv(g, h, weight=ext_w)
-    assert h.shape == (ndst, 2)
+        h_out = conv(g, h, weight=ext_w)
+    assert h_out.shape == (ndst, 2)
+
+    if not isinstance(g, dgl.DGLGraph) and len(g.ntypes) == 2:
+        # bipartite, should also accept pair of tensors
+        if weight:
+            h_out2 = conv(g, (h, h_dst))
+        else:
+            h_out2 = conv(g, (h, h_dst), weight=ext_w)
+        assert h_out2.shape == (ndst, 2)
+        assert F.array_equal(h_out, h_out2)
 
 def test_simple_pool():
     ctx = F.ctx()