Skip to content

GitLab

Unverified Commit 701b4fcc authored by Quan (Andy) Gan's avatar Quan (Andy) Gan Committed by GitHub
Browse files

[Sampling] New sampling pipeline plus asynchronous prefetching (#3665) 

* initial update

* more

* more

* multi-gpu example

* cluster gcn, finalize homogeneous

* more explanation

* fix

* bunch of fixes

* fix

* RGAT example and more fixes

* shadow-gnn sampler and some changes in unit test

* fix

* wth

* more fixes

* remove shadow+node/edge dataloader tests for possible ux changes

* lints

* add legacy dataloading import just in case

* fix

* update pylint for f-strings

* fix

* lint

* lint

* lint again

* cherry-picking commit fa9f494

* oops

* fix

* add sample_neighbors in dist_graph

* fix

* lint

* fix

* fix

* fix

* fix tutorial

* fix

* fix

* fix

* fix warning

* remove debug

* add get_foo_storage apis

* lint
parent 5152a879
Hide whitespace changes
Inline Side-by-side
Showing with 2651 additions and 1108 deletions
python/dgl/dataloading/async_transferer.py → python/dgl/_dataloading/async_transferer.py
View file @ 701b4fcc
......@@ -101,4 +101,4 @@ class AsyncTransferer(object):
return Transfer(transfer_id, self._handle)
_init_api("dgl.dataloading.async_transferer")
_init_api("dataloading.async_transferer", "dgl._dataloading.async_transferer")
python/dgl/_dataloading/cluster_gcn.py 0 → 100644
View file @ 701b4fcc
"""Cluster-GCN subgraph iterators."""
import os
import pickle
import numpy as np
from ..transform import metis_partition_assignment
from .. import backend as F
from .dataloader import SubgraphIterator
class ClusterGCNSubgraphIterator(SubgraphIterator):
"""Subgraph sampler following that of ClusterGCN.
This sampler first partitions the graph with METIS partitioning, then it caches the nodes of
each partition to a file within the given cache directory.
This is used in conjunction with :class:`dgl.dataloading.pytorch.GraphDataLoader`.
Notes
-----
The graph must be homogeneous and on CPU.
Parameters
----------
g : DGLGraph
The original graph.
num_partitions : int
The number of partitions.
cache_directory : str
The path to the cache directory for storing the partition result.
refresh : bool
If True, recompute the partition.
Examples
--------
Assuming that you have a graph ``g``:
>>> sgiter = dgl.dataloading.ClusterGCNSubgraphIterator(
... g, num_partitions=100, cache_directory='.', refresh=True)
>>> dataloader = dgl.dataloading.GraphDataLoader(sgiter, batch_size=4, num_workers=0)
>>> for subgraph_batch in dataloader:
... train_on(subgraph_batch)
"""
def __init__(self, g, num_partitions, cache_directory, refresh=False):
if os.name == 'nt':
raise NotImplementedError("METIS partitioning is not supported on Windows yet.")
super().__init__(g)
# First see if the cache is already there. If so, directly read from cache.
if not refresh and self._load_parts(cache_directory):
return
# Otherwise, build the cache.
assignment = F.asnumpy(metis_partition_assignment(g, num_partitions))
self._save_parts(assignment, cache_directory)
def _cache_file_path(self, cache_directory):
return os.path.join(cache_directory, 'cluster_gcn_cache')
def _load_parts(self, cache_directory):
path = self._cache_file_path(cache_directory)
if not os.path.exists(path):
return False
with open(path, 'rb') as file_:
self.part_indptr, self.part_indices = pickle.load(file_)
return True
def _save_parts(self, assignment, cache_directory):
os.makedirs(cache_directory, exist_ok=True)
self.part_indices = np.argsort(assignment)
num_nodes_per_part = np.bincount(assignment)
self.part_indptr = np.insert(np.cumsum(num_nodes_per_part), 0, 0)
with open(self._cache_file_path(cache_directory), 'wb') as file_:
pickle.dump((self.part_indptr, self.part_indices), file_)
def __len__(self):
return self.part_indptr.shape[0] - 1
def __getitem__(self, i):
nodes = self.part_indices[self.part_indptr[i]:self.part_indptr[i+1]]
return self.g.subgraph(nodes)
python/dgl/_dataloading/dataloader.py 0 → 100644
View file @ 701b4fcc
"""Data loaders"""
from collections.abc import Mapping, Sequence
from abc import ABC, abstractproperty, abstractmethod
import re
import numpy as np
from .. import transform
from ..base import NID, EID
from .. import backend as F
from .. import utils
from ..batch import batch
from ..convert import heterograph
from ..heterograph import DGLHeteroGraph as DGLGraph
from ..distributed.dist_graph import DistGraph
from ..utils import to_device
def _tensor_or_dict_to_numpy(ids):
if isinstance(ids, Mapping):
return {k: F.zerocopy_to_numpy(v) for k, v in ids.items()}
else:
return F.zerocopy_to_numpy(ids)
def _locate_eids_to_exclude(frontier_parent_eids, exclude_eids):
"""Find the edges whose IDs in parent graph appeared in exclude_eids.
Note that both arguments are numpy arrays or numpy dicts.
"""
if isinstance(frontier_parent_eids, Mapping):
result = {
k: np.isin(frontier_parent_eids[k], exclude_eids[k]).nonzero()[0]
for k in frontier_parent_eids.keys() if k in exclude_eids.keys()}
return {k: F.zerocopy_from_numpy(v) for k, v in result.items()}
else:
result = np.isin(frontier_parent_eids, exclude_eids).nonzero()[0]
return F.zerocopy_from_numpy(result)
class _EidExcluder():
def __init__(self, exclude_eids):
device = None
if isinstance(exclude_eids, Mapping):
for _, v in exclude_eids.items():
if device is None:
device = F.context(v)
break
else:
device = F.context(exclude_eids)
self._exclude_eids = None
self._filter = None
if device == F.cpu():
# TODO(nv-dlasalle): Once Filter is implemented for the CPU, we
# should just use that irregardless of the device.
self._exclude_eids = (
_tensor_or_dict_to_numpy(exclude_eids) if exclude_eids is not None else None)
else:
if isinstance(exclude_eids, Mapping):
self._filter = {k: utils.Filter(v) for k, v in exclude_eids.items()}
else:
self._filter = utils.Filter(exclude_eids)
def _find_indices(self, parent_eids):
""" Find the set of edge indices to remove.
"""
if self._exclude_eids is not None:
parent_eids_np = _tensor_or_dict_to_numpy(parent_eids)
return _locate_eids_to_exclude(parent_eids_np, self._exclude_eids)
else:
assert self._filter is not None
if isinstance(parent_eids, Mapping):
located_eids = {k: self._filter[k].find_included_indices(parent_eids[k])
for k, v in parent_eids.items() if k in self._filter}
else:
located_eids = self._filter.find_included_indices(parent_eids)
return located_eids
def __call__(self, frontier):
parent_eids = frontier.edata[EID]
located_eids = self._find_indices(parent_eids)
if not isinstance(located_eids, Mapping):
# (BarclayII) If frontier already has a EID field and located_eids is empty,
# the returned graph will keep EID intact. Otherwise, EID will change
# to the mapping from the new graph to the old frontier.
# So we need to test if located_eids is empty, and do the remapping ourselves.
if len(located_eids) > 0:
frontier = transform.remove_edges(
frontier, located_eids, store_ids=True)
frontier.edata[EID] = F.gather_row(parent_eids, frontier.edata[EID])
else:
# (BarclayII) remove_edges only accepts removing one type of edges,
# so I need to keep track of the edge IDs left one by one.
new_eids = parent_eids.copy()
for k, v in located_eids.items():
if len(v) > 0:
frontier = transform.remove_edges(
frontier, v, etype=k, store_ids=True)
new_eids[k] = F.gather_row(parent_eids[k], frontier.edges[k].data[EID])
frontier.edata[EID] = new_eids
return frontier
def exclude_edges(subg, exclude_eids, device):
"""Find and remove from the subgraph the edges whose IDs in the parent
graph are given.
Parameters
----------
subg : DGLGraph
The subgraph. Must have ``dgl.EID`` field containing the original
edge IDs in the parent graph.
exclude_eids : Tensor or dict
The edge IDs to exclude.
device : device
The output device of the graph.
Returns
-------
DGLGraph
The new subgraph with edges removed. The ``dgl.EID`` field contains
the original edge IDs in the same parent graph.
"""
if exclude_eids is None:
return subg
if device is not None:
if isinstance(exclude_eids, Mapping):
exclude_eids = {k: F.copy_to(v, device) \
for k, v in exclude_eids.items()}
else:
exclude_eids = F.copy_to(exclude_eids, device)
excluder = _EidExcluder(exclude_eids)
return subg if excluder is None else excluder(subg)
def _find_exclude_eids_with_reverse_id(g, eids, reverse_eid_map):
if isinstance(eids, Mapping):
eids = {g.to_canonical_etype(k): v for k, v in eids.items()}
exclude_eids = {
k: F.cat([v, F.gather_row(reverse_eid_map[k], v)], 0)
for k, v in eids.items()}
else:
exclude_eids = F.cat([eids, F.gather_row(reverse_eid_map, eids)], 0)
return exclude_eids
def _find_exclude_eids_with_reverse_types(g, eids, reverse_etype_map):
exclude_eids = {g.to_canonical_etype(k): v for k, v in eids.items()}
reverse_etype_map = {
g.to_canonical_etype(k): g.to_canonical_etype(v)
for k, v in reverse_etype_map.items()}
exclude_eids.update({reverse_etype_map[k]: v for k, v in exclude_eids.items()})
return exclude_eids
def _find_exclude_eids(g, exclude_mode, eids, **kwargs):
"""Find all edge IDs to exclude according to :attr:`exclude_mode`.
Parameters
----------
g : DGLGraph
The graph.
exclude_mode : str, optional
Can be either of the following,
None (default)
Does not exclude any edge.
'self'
Exclude the given edges themselves but nothing else.
'reverse_id'
Exclude all edges specified in ``eids``, as well as their reverse edges
of the same edge type.
The mapping from each edge ID to its reverse edge ID is specified in
the keyword argument ``reverse_eid_map``.
This mode assumes that the reverse of an edge with ID ``e`` and type
``etype`` will have ID ``reverse_eid_map[e]`` and type ``etype``.
'reverse_types'
Exclude all edges specified in ``eids``, as well as their reverse
edges of the corresponding edge types.
The mapping from each edge type to its reverse edge type is specified
in the keyword argument ``reverse_etype_map``.
This mode assumes that the reverse of an edge with ID ``e`` and type ``etype``
will have ID ``e`` and type ``reverse_etype_map[etype]``.
eids : Tensor or dict[etype, Tensor]
The edge IDs.
reverse_eid_map : Tensor or dict[etype, Tensor]
The mapping from edge ID to its reverse edge ID.
reverse_etype_map : dict[etype, etype]
The mapping from edge etype to its reverse edge type.
"""
if exclude_mode is None:
return None
elif exclude_mode == 'self':
return eids
elif exclude_mode == 'reverse_id':
return _find_exclude_eids_with_reverse_id(g, eids, kwargs['reverse_eid_map'])
elif exclude_mode == 'reverse_types':
return _find_exclude_eids_with_reverse_types(g, eids, kwargs['reverse_etype_map'])
else:
raise ValueError('unsupported mode {}'.format(exclude_mode))
class Sampler(object):
"""An abstract class that takes in a graph and a set of seed nodes and returns a
structure representing a smaller portion of the graph for computation. It can
be either a list of bipartite graphs (i.e. :class:`BlockSampler`), or a single
subgraph.
"""
def __init__(self, output_ctx=None):
self.set_output_context(output_ctx)
def sample(self, g, seed_nodes, exclude_eids=None):
"""Sample a structure from the graph.
Parameters
----------
g : DGLGraph
The original graph.
seed_nodes : Tensor or dict[ntype, Tensor]
The destination nodes by type.
If the graph only has one node type, one can just specify a single tensor
of node IDs.
exclude_eids : Tensor or dict[etype, Tensor]
The edges to exclude from computation dependency.
Returns
-------
Tensor or dict[ntype, Tensor]
The nodes whose input features are required for computing the output
representation of :attr:`seed_nodes`.
any
Any data representing the structure.
"""
raise NotImplementedError
def set_output_context(self, ctx):
"""Set the device the generated block or subgraph will be output to.
This should only be set to a cuda device, when multi-processing is not
used in the dataloader (e.g., num_workers is 0).
Parameters
----------
ctx : DGLContext, default None
The device context the sampled blocks will be stored on. This
should only be a CUDA context if multiprocessing is not used in
the dataloader (e.g., num_workers is 0). If this is None, the
sampled blocks will be stored on the same device as the input
graph.
"""
if ctx is not None:
self.output_device = F.to_backend_ctx(ctx)
else:
self.output_device = None
class BlockSampler(Sampler):
"""Abstract class specifying the neighborhood sampling strategy for DGL data loaders.
The main method for BlockSampler is :meth:`sample`,
which generates a list of message flow graphs (MFGs) for a multi-layer GNN given a set of
seed nodes to have their outputs computed.
The default implementation of :meth:`sample` is
to repeat :attr:`num_layers` times the following procedure from the last layer to the first
layer:
* Obtain a frontier. The frontier is defined as a graph with the same nodes as the
original graph but only the edges involved in message passing on the current layer.
Customizable via :meth:`sample_frontier`.
* Optionally, if the task is link prediction or edge classfication, remove edges
connecting training node pairs. If the graph is undirected, also remove the
reverse edges. This is controlled by the argument :attr:`exclude_eids` in
:meth:`sample` method.
* Convert the frontier into a MFG.
* Optionally assign the IDs of the edges in the original graph selected in the first step
to the MFG, controlled by the argument ``return_eids`` in
:meth:`sample` method.
* Prepend the MFG to the MFG list to be returned.
All subclasses should override :meth:`sample_frontier`
method while specifying the number of layers to sample in :attr:`num_layers` argument.
Parameters
----------
num_layers : int
The number of layers to sample.
return_eids : bool, default False
Whether to return the edge IDs involved in message passing in the MFG.
If True, the edge IDs will be stored as an edge feature named ``dgl.EID``.
output_ctx : DGLContext, default None
The context the sampled blocks will be stored on. This should only be
a CUDA context if multiprocessing is not used in the dataloader (e.g.,
num_workers is 0). If this is None, the sampled blocks will be stored
on the same device as the input graph.
exclude_edges_in_frontier : bool, default False
If True, the :func:`sample_frontier` method will receive an argument
:attr:`exclude_eids` containing the edge IDs from the original graph to exclude.
The :func:`sample_frontier` method must return a graph that does not contain
the edges corresponding to the excluded edges. No additional postprocessing
will be done.
Otherwise, the edges will be removed *after* :func:`sample_frontier` returns.
Notes
-----
For the concept of frontiers and MFGs, please refer to
:ref:`User Guide Section 6 <guide-minibatch>` and
:doc:`Minibatch Training Tutorials <tutorials/large/L0_neighbor_sampling_overview>`.
"""
def __init__(self, num_layers, return_eids=False, output_ctx=None):
super().__init__(output_ctx)
self.num_layers = num_layers
self.return_eids = return_eids
# pylint: disable=unused-argument
@staticmethod
def assign_block_eids(block, frontier):
"""Assigns edge IDs from the original graph to the message flow graph (MFG).
See also
--------
BlockSampler
"""
for etype in block.canonical_etypes:
block.edges[etype].data[EID] = frontier.edges[etype].data[EID][
block.edges[etype].data[EID]]
return block
# This is really a hack working around the lack of GPU-based neighbor sampling
# with edge exclusion.
@classmethod
def exclude_edges_in_frontier(cls, g):
"""Returns whether the sampler will exclude edges in :func:`sample_frontier`.
If this method returns True, the method :func:`sample_frontier` will receive an
argument :attr:`exclude_eids` from :func:`sample`. :func:`sample_frontier`
is then responsible for removing those edges.
If this method returns False, :func:`sample` will be responsible for
removing the edges.
When subclassing :class:`BlockSampler`, this method should return True when you
would like to remove the excluded edges in your :func:`sample_frontier` method.
By default this method returns False.
Parameters
----------
g : DGLGraph
The original graph
Returns
-------
bool
Whether :func:`sample_frontier` will receive an argument :attr:`exclude_eids`.
"""
return False
def sample_frontier(self, block_id, g, seed_nodes, exclude_eids=None):
"""Generate the frontier given the destination nodes.
The subclasses should override this function.
Parameters
----------
block_id : int
Represents which GNN layer the frontier is generated for.
g : DGLGraph
The original graph.
seed_nodes : Tensor or dict[ntype, Tensor]
The destination nodes by node type.
If the graph only has one node type, one can just specify a single tensor
of node IDs.
exclude_eids: Tensor or dict
Edge IDs to exclude during sampling neighbors for the seed nodes.
This argument can take a single ID tensor or a dictionary of edge types and ID tensors.
If a single tensor is given, the graph must only have one type of nodes.
Returns
-------
DGLGraph
The frontier generated for the current layer.
Notes
-----
For the concept of frontiers and MFGs, please refer to
:ref:`User Guide Section 6 <guide-minibatch>` and
:doc:`Minibatch Training Tutorials <tutorials/large/L0_neighbor_sampling_overview>`.
"""
raise NotImplementedError
def sample(self, g, seed_nodes, exclude_eids=None):
"""Generate the a list of MFGs given the destination nodes.
Parameters
----------
g : DGLGraph
The original graph.
seed_nodes : Tensor or dict[ntype, Tensor]
The destination nodes by node type.
If the graph only has one node type, one can just specify a single tensor
of node IDs.
exclude_eids : Tensor or dict[etype, Tensor]
The edges to exclude from computation dependency.
Returns
-------
list[DGLGraph]
The MFGs generated for computing the multi-layer GNN output.
Notes
-----
For the concept of frontiers and MFGs, please refer to
:ref:`User Guide Section 6 <guide-minibatch>` and
:doc:`Minibatch Training Tutorials <tutorials/large/L0_neighbor_sampling_overview>`.
"""
blocks = []
if isinstance(g, DistGraph):
# TODO:(nv-dlasalle) dist graphs may not have an associated graph,
# causing an error when trying to fetch the device, so for now,
# always assume the distributed graph's device is CPU.
graph_device = F.cpu()
else:
graph_device = g.device
for block_id in reversed(range(self.num_layers)):
seed_nodes_in = to_device(seed_nodes, graph_device)
if self.exclude_edges_in_frontier(g):
frontier = self.sample_frontier(
block_id, g, seed_nodes_in, exclude_eids=exclude_eids)
else:
frontier = self.sample_frontier(block_id, g, seed_nodes_in)
if self.output_device is not None:
frontier = frontier.to(self.output_device)
seed_nodes_out = to_device(seed_nodes, self.output_device)
else:
seed_nodes_out = seed_nodes
# Removing edges from the frontier for link prediction training falls
# into the category of frontier postprocessing
if not self.exclude_edges_in_frontier(g):
frontier = exclude_edges(frontier, exclude_eids, self.output_device)
block = transform.to_block(frontier, seed_nodes_out)
if self.return_eids:
self.assign_block_eids(block, frontier)
seed_nodes = {ntype: block.srcnodes[ntype].data[NID] for ntype in block.srctypes}
blocks.insert(0, block)
return blocks[0].srcdata[NID], blocks[-1].dstdata[NID], blocks
def sample_blocks(self, g, seed_nodes, exclude_eids=None):
"""Deprecated and identical to :meth:`sample`.
"""
return self.sample(g, seed_nodes, exclude_eids)
class Collator(ABC):
"""Abstract DGL collator for training GNNs on downstream tasks stochastically.
Provides a :attr:`dataset` object containing the collection of all nodes or edges,
as well as a :attr:`collate` method that combines a set of items from
:attr:`dataset` and obtains the message flow graphs (MFGs).
Notes
-----
For the concept of MFGs, please refer to
:ref:`User Guide Section 6 <guide-minibatch>` and
:doc:`Minibatch Training Tutorials <tutorials/large/L0_neighbor_sampling_overview>`.
"""
@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 MFGs.
Parameters
----------
items : list[str, int]
The list of node or edge IDs or type-ID pairs.
Notes
-----
For the concept of MFGs, please refer to
:ref:`User Guide Section 6 <guide-minibatch>` and
:doc:`Minibatch Training Tutorials <tutorials/large/L0_neighbor_sampling_overview>`.
"""
raise NotImplementedError
class NodeCollator(Collator):
"""DGL collator to combine nodes and their computation dependencies within a minibatch for
training node classification or regression on a single graph with neighborhood sampling.
Parameters
----------
g : DGLGraph
The graph.
nids : Tensor or dict[ntype, Tensor]
The node set to compute outputs.
graph_sampler : dgl.dataloading.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.dataloading.MultiLayerNeighborSampler([15, 10, 5])
>>> collator = dgl.dataloading.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)
Notes
-----
For the concept of MFGs, please refer to
:ref:`User Guide Section 6 <guide-minibatch>` and
:doc:`Minibatch Training Tutorials <tutorials/large/L0_neighbor_sampling_overview>`.
"""
def __init__(self, g, nids, graph_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.graph_sampler = graph_sampler
self.nids = utils.prepare_tensor_or_dict(g, nids, 'nids')
self._dataset = utils.maybe_flatten_dict(self.nids)
@property
def dataset(self):
return self._dataset
def collate(self, items):
"""Find the list of MFGs necessary for computing the representation of given
nodes for a node classification/regression task.
Parameters
----------
items : list[int] or list[tuple[str, int]]
Either a list of node IDs (for homogeneous graphs), or a list of node type-ID
pairs (for heterogeneous graphs).
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.
MFGs : list[DGLGraph]
The list of MFGs 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)
items = utils.prepare_tensor_or_dict(self.g, items, 'items')
input_nodes, output_nodes, blocks = self.graph_sampler.sample_blocks(self.g, items)
return input_nodes, output_nodes, blocks
class EdgeCollator(Collator):
"""DGL collator to combine edges and their computation dependencies within a minibatch for
training edge classification, edge regression, or link prediction on a single graph
with neighborhood sampling.
Given a set of edges, the collate function will yield
* A tensor of input nodes necessary for computing the representation on edges, or
a dictionary of node type names and such tensors.
* A subgraph that contains only the edges in the minibatch and their incident nodes.
Note that the graph has an identical metagraph with the original graph.
* If a negative sampler is given, another graph that contains the "negative edges",
connecting the source and destination nodes yielded from the given negative sampler.
* A list of MFGs necessary for computing the representation of the incident nodes
of the edges in the minibatch.
Parameters
----------
g : DGLGraph
The graph from which the edges are iterated in minibatches and the subgraphs
are generated.
eids : Tensor or dict[etype, Tensor]
The edge set in graph :attr:`g` to compute outputs.
graph_sampler : dgl.dataloading.BlockSampler
The neighborhood sampler.
g_sampling : DGLGraph, optional
The graph where neighborhood sampling and message passing is performed.
Note that this is not necessarily the same as :attr:`g`.
If None, assume to be the same as :attr:`g`.
exclude : str, optional
Whether and how to exclude dependencies related to the sampled edges in the
minibatch. Possible values are
* None, which excludes nothing.
* ``'self'``, which excludes the sampled edges themselves but nothing else.
* ``'reverse_id'``, which excludes the reverse edges of the sampled edges. The said
reverse edges have the same edge type as the sampled edges. Only works
on edge types whose source node type is the same as its destination node type.
* ``'reverse_types'``, which excludes the reverse edges of the sampled edges. The
said reverse edges have different edge types from the sampled edges.
If ``g_sampling`` is given, ``exclude`` is ignored and will be always ``None``.
reverse_eids : Tensor or dict[etype, Tensor], optional
A tensor of reverse edge ID mapping. The i-th element indicates the ID of
the i-th edge's reverse edge.
If the graph is heterogeneous, this argument requires a dictionary of edge
types and the reverse edge ID mapping tensors.
Required and only used when ``exclude`` is set to ``reverse_id``.
For heterogeneous graph this will be a dict of edge type and edge IDs. Note that
only the edge types whose source node type is the same as destination node type
are needed.
reverse_etypes : dict[etype, etype], optional
The mapping from the edge type to its reverse edge type.
Required and only used when ``exclude`` is set to ``reverse_types``.
negative_sampler : callable, optional
The negative sampler. Can be omitted if no negative sampling is needed.
The negative sampler must be a callable that takes in the following arguments:
* The original (heterogeneous) graph.
* The ID array of sampled edges in the minibatch, or the dictionary of edge
types and ID array of sampled edges in the minibatch if the graph is
heterogeneous.
It should return
* A pair of source and destination node ID arrays as negative samples,
or a dictionary of edge types and such pairs if the graph is heterogenenous.
A set of builtin negative samplers are provided in
:ref:`the negative sampling module <api-dataloading-negative-sampling>`.
Examples
--------
The following example shows how to train a 3-layer GNN for edge classification on a
set of edges ``train_eid`` on a homogeneous undirected graph. Each node takes
messages from all neighbors.
Say that you have an array of source node IDs ``src`` and another array of destination
node IDs ``dst``. One can make it bidirectional by adding another set of edges
that connects from ``dst`` to ``src``:
>>> g = dgl.graph((torch.cat([src, dst]), torch.cat([dst, src])))
One can then know that the ID difference of an edge and its reverse edge is ``|E|``,
where ``|E|`` is the length of your source/destination array. The reverse edge
mapping can be obtained by
>>> E = len(src)
>>> reverse_eids = torch.cat([torch.arange(E, 2 * E), torch.arange(0, E)])
Note that the sampled edges as well as their reverse edges are removed from
computation dependencies of the incident nodes. This is a common trick to avoid
information leakage.
>>> sampler = dgl.dataloading.MultiLayerNeighborSampler([15, 10, 5])
>>> collator = dgl.dataloading.EdgeCollator(
... g, train_eid, sampler, exclude='reverse_id',
... reverse_eids=reverse_eids)
>>> 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, pair_graph, blocks in dataloader:
... train_on(input_nodes, pair_graph, blocks)
To train a 3-layer GNN for link prediction on a set of edges ``train_eid`` on a
homogeneous graph where each node takes messages from all neighbors (assume the
backend is PyTorch), with 5 uniformly chosen negative samples per edge:
>>> sampler = dgl.dataloading.MultiLayerNeighborSampler([15, 10, 5])
>>> neg_sampler = dgl.dataloading.negative_sampler.Uniform(5)
>>> collator = dgl.dataloading.EdgeCollator(
... g, train_eid, sampler, exclude='reverse_id',
... reverse_eids=reverse_eids, negative_sampler=neg_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, pos_pair_graph, neg_pair_graph, blocks in dataloader:
... train_on(input_nodse, pair_graph, neg_pair_graph, blocks)
For heterogeneous graphs, the reverse of an edge may have a different edge type
from the original edge. For instance, consider that you have an array of
user-item clicks, representated by a user array ``user`` and an item array ``item``.
You may want to build a heterogeneous graph with a user-click-item relation and an
item-clicked-by-user relation.
>>> g = dgl.heterograph({
... ('user', 'click', 'item'): (user, item),
... ('item', 'clicked-by', 'user'): (item, user)})
To train a 3-layer GNN for edge classification on a set of edges ``train_eid`` with
type ``click``, you can write
>>> sampler = dgl.dataloading.MultiLayerNeighborSampler([15, 10, 5])
>>> collator = dgl.dataloading.EdgeCollator(
... g, {'click': train_eid}, sampler, exclude='reverse_types',
... reverse_etypes={'click': 'clicked-by', 'clicked-by': 'click'})
>>> 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, pair_graph, blocks in dataloader:
... train_on(input_nodes, pair_graph, blocks)
To train a 3-layer GNN for link prediction on a set of edges ``train_eid`` with type
``click``, you can write
>>> sampler = dgl.dataloading.MultiLayerNeighborSampler([15, 10, 5])
>>> neg_sampler = dgl.dataloading.negative_sampler.Uniform(5)
>>> collator = dgl.dataloading.EdgeCollator(
... g, train_eid, sampler, exclude='reverse_types',
... reverse_etypes={'click': 'clicked-by', 'clicked-by': 'click'},
... negative_sampler=neg_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, pos_pair_graph, neg_pair_graph, blocks in dataloader:
... train_on(input_nodes, pair_graph, neg_pair_graph, blocks)
Notes
-----
For the concept of MFGs, please refer to
:ref:`User Guide Section 6 <guide-minibatch>` and
:doc:`Minibatch Training Tutorials <tutorials/large/L0_neighbor_sampling_overview>`.
"""
def __init__(self, g, eids, graph_sampler, g_sampling=None, exclude=None,
reverse_eids=None, reverse_etypes=None, negative_sampler=None):
self.g = g
if not isinstance(eids, Mapping):
assert len(g.etypes) == 1, \
"eids should be a dict of etype and ids for graph with multiple etypes"
self.graph_sampler = graph_sampler
# One may wish to iterate over the edges in one graph while perform sampling in
# another graph. This may be the case for iterating over validation and test
# edge set while perform neighborhood sampling on the graph formed by only
# the training edge set.
# See GCMC for an example usage.
if g_sampling is not None:
self.g_sampling = g_sampling
self.exclude = None
else:
self.g_sampling = self.g
self.exclude = exclude
self.reverse_eids = reverse_eids
self.reverse_etypes = reverse_etypes
self.negative_sampler = negative_sampler
self.eids = utils.prepare_tensor_or_dict(g, eids, 'eids')
self._dataset = utils.maybe_flatten_dict(self.eids)
@property
def dataset(self):
return self._dataset
def _collate(self, items):
if isinstance(items[0], tuple):
# returns a list of pairs: group them by node types into a dict
items = utils.group_as_dict(items)
items = utils.prepare_tensor_or_dict(self.g_sampling, items, 'items')
pair_graph = self.g.edge_subgraph(items)
seed_nodes = pair_graph.ndata[NID]
exclude_eids = _find_exclude_eids(
self.g_sampling,
self.exclude,
items,
reverse_eid_map=self.reverse_eids,
reverse_etype_map=self.reverse_etypes)
input_nodes, _, blocks = self.graph_sampler.sample_blocks(
self.g_sampling, seed_nodes, exclude_eids=exclude_eids)
return input_nodes, pair_graph, blocks
def _collate_with_negative_sampling(self, items):
if isinstance(items[0], tuple):
# returns a list of pairs: group them by node types into a dict
items = utils.group_as_dict(items)
items = utils.prepare_tensor_or_dict(self.g_sampling, items, 'items')
pair_graph = self.g.edge_subgraph(items, relabel_nodes=False)
induced_edges = pair_graph.edata[EID]
neg_srcdst = self.negative_sampler(self.g, items)
if not isinstance(neg_srcdst, Mapping):
assert len(self.g.etypes) == 1, \
'graph has multiple or no edge types; '\
'please return a dict in negative sampler.'
neg_srcdst = {self.g.canonical_etypes[0]: neg_srcdst}
# Get dtype from a tuple of tensors
dtype = F.dtype(list(neg_srcdst.values())[0][0])
ctx = F.context(pair_graph)
neg_edges = {
etype: neg_srcdst.get(etype, (F.copy_to(F.tensor([], dtype), ctx),
F.copy_to(F.tensor([], dtype), ctx)))
for etype in self.g.canonical_etypes}
neg_pair_graph = heterograph(
neg_edges, {ntype: self.g.number_of_nodes(ntype) for ntype in self.g.ntypes})
pair_graph, neg_pair_graph = transform.compact_graphs([pair_graph, neg_pair_graph])
pair_graph.edata[EID] = induced_edges
seed_nodes = pair_graph.ndata[NID]
exclude_eids = _find_exclude_eids(
self.g_sampling,
self.exclude,
items,
reverse_eid_map=self.reverse_eids,
reverse_etype_map=self.reverse_etypes)
input_nodes, _, blocks = self.graph_sampler.sample_blocks(
self.g_sampling, seed_nodes, exclude_eids=exclude_eids)
return input_nodes, pair_graph, neg_pair_graph, blocks
def collate(self, items):
"""Combines the sampled edges into a minibatch for edge classification, edge
regression, and link prediction tasks.
Parameters
----------
items : list[int] or list[tuple[str, int]]
Either a list of edge IDs (for homogeneous graphs), or a list of edge type-ID
pairs (for heterogeneous graphs).
Returns
-------
Either ``(input_nodes, pair_graph, blocks)``, or
``(input_nodes, pair_graph, negative_pair_graph, blocks)`` if negative sampling is
enabled.
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.
pair_graph : DGLGraph
The graph that contains only the edges in the minibatch as well as their incident
nodes.
Note that the metagraph of this graph will be identical to that of the original
graph.
negative_pair_graph : DGLGraph
The graph that contains only the edges connecting the source and destination nodes
yielded from the given negative sampler, if negative sampling is enabled.
Note that the metagraph of this graph will be identical to that of the original
graph.
blocks : list[DGLGraph]
The list of MFGs necessary for computing the representation of the edges.
"""
if self.negative_sampler is None:
return self._collate(items)
else:
return self._collate_with_negative_sampling(items)
class GraphCollator(object):
"""Given a set of graphs as well as their graph-level data, the collate function will batch the
graphs into a batched graph, and stack the tensors into a single bigger tensor. If the
example is a container (such as sequences or mapping), the collate function preserves
the structure and collates each of the elements recursively.
If the set of graphs has no graph-level data, the collate function will yield a batched graph.
Examples
--------
To train a GNN for graph classification on a set of graphs in ``dataset`` (assume
the backend is PyTorch):
>>> dataloader = dgl.dataloading.GraphDataLoader(
... dataset, batch_size=1024, shuffle=True, drop_last=False, num_workers=4)
>>> for batched_graph, labels in dataloader:
... train_on(batched_graph, labels)
"""
def __init__(self):
self.graph_collate_err_msg_format = (
"graph_collate: batch must contain DGLGraph, tensors, numpy arrays, "
"numbers, dicts or lists; found {}")
self.np_str_obj_array_pattern = re.compile(r'[SaUO]')
#This implementation is based on torch.utils.data._utils.collate.default_collate
def collate(self, items):
"""This function is similar to ``torch.utils.data._utils.collate.default_collate``.
It combines the sampled graphs and corresponding graph-level data
into a batched graph and tensors.
Parameters
----------
items : list of data points or tuples
Elements in the list are expected to have the same length.
Each sub-element will be batched as a batched graph, or a
batched tensor correspondingly.
Returns
-------
A tuple of the batching results.
"""
elem = items[0]
elem_type = type(elem)
if isinstance(elem, DGLGraph):
batched_graphs = batch(items)
return batched_graphs
elif F.is_tensor(elem):
return F.stack(items, 0)
elif elem_type.__module__ == 'numpy' and elem_type.__name__ != 'str_' \
and elem_type.__name__ != 'string_':
if elem_type.__name__ == 'ndarray' or elem_type.__name__ == 'memmap':
# array of string classes and object
if self.np_str_obj_array_pattern.search(elem.dtype.str) is not None:
raise TypeError(self.graph_collate_err_msg_format.format(elem.dtype))
return self.collate([F.tensor(b) for b in items])
elif elem.shape == (): # scalars
return F.tensor(items)
elif isinstance(elem, float):
return F.tensor(items, dtype=F.float64)
elif isinstance(elem, int):
return F.tensor(items)
elif isinstance(elem, (str, bytes)):
return items
elif isinstance(elem, Mapping):
return {key: self.collate([d[key] for d in items]) for key in elem}
elif isinstance(elem, tuple) and hasattr(elem, '_fields'): # namedtuple
return elem_type(*(self.collate(samples) for samples in zip(*items)))
elif isinstance(elem, Sequence):
# check to make sure that the elements in batch have consistent size
item_iter = iter(items)
elem_size = len(next(item_iter))
if not all(len(elem) == elem_size for elem in item_iter):
raise RuntimeError('each element in list of batch should be of equal size')
transposed = zip(*items)
return [self.collate(samples) for samples in transposed]
raise TypeError(self.graph_collate_err_msg_format.format(elem_type))
class SubgraphIterator(object):
"""Abstract class representing an iterator that yields a subgraph given a graph.
"""
def __init__(self, g):
self.g = g
python/dgl/_dataloading/negative_sampler.py 0 → 100644
View file @ 701b4fcc
"""Negative samplers"""
from collections.abc import Mapping
from .. import backend as F
from ..sampling import global_uniform_negative_sampling
class _BaseNegativeSampler(object):
def _generate(self, g, eids, canonical_etype):
raise NotImplementedError
def __call__(self, g, eids):
"""Returns negative samples.
Parameters
----------
g : DGLGraph
The graph.
eids : Tensor or dict[etype, Tensor]
The sampled edges in the minibatch.
Returns
-------
tuple[Tensor, Tensor] or dict[etype, tuple[Tensor, Tensor]]
The returned source-destination pairs as negative samples.
"""
if isinstance(eids, Mapping):
eids = {g.to_canonical_etype(k): v for k, v in eids.items()}
neg_pair = {k: self._generate(g, v, k) for k, v in eids.items()}
else:
assert len(g.etypes) == 1, \
'please specify a dict of etypes and ids for graphs with multiple edge types'
neg_pair = self._generate(g, eids, g.canonical_etypes[0])
return neg_pair
class PerSourceUniform(_BaseNegativeSampler):
"""Negative sampler that randomly chooses negative destination nodes
for each source node according to a uniform distribution.
For each edge ``(u, v)`` of type ``(srctype, etype, dsttype)``, DGL generates
:attr:`k` pairs of negative edges ``(u, v')``, where ``v'`` is chosen
uniformly from all the nodes of type ``dsttype``. The resulting edges will
also have type ``(srctype, etype, dsttype)``.
Parameters
----------
k : int
The number of negative samples per edge.
Examples
--------
>>> g = dgl.graph(([0, 1, 2], [1, 2, 3]))
>>> neg_sampler = dgl.dataloading.negative_sampler.PerSourceUniform(2)
>>> neg_sampler(g, torch.tensor([0, 1]))
(tensor([0, 0, 1, 1]), tensor([1, 0, 2, 3]))
"""
def __init__(self, k):
self.k = k
def _generate(self, g, eids, canonical_etype):
_, _, vtype = canonical_etype
shape = F.shape(eids)
dtype = F.dtype(eids)
ctx = F.context(eids)
shape = (shape[0] * self.k,)
src, _ = g.find_edges(eids, etype=canonical_etype)
src = F.repeat(src, self.k, 0)
dst = F.randint(shape, dtype, ctx, 0, g.number_of_nodes(vtype))
return src, dst
# Alias
Uniform = PerSourceUniform
class GlobalUniform(_BaseNegativeSampler):
"""Negative sampler that randomly chooses negative source-destination pairs according
to a uniform distribution.
For each edge ``(u, v)`` of type ``(srctype, etype, dsttype)``, DGL generates at most
:attr:`k` pairs of negative edges ``(u', v')``, where ``u'`` is chosen uniformly from
all the nodes of type ``srctype`` and ``v'`` is chosen uniformly from all the nodes
of type ``dsttype``. The resulting edges will also have type
``(srctype, etype, dsttype)``. DGL guarantees that the sampled pairs will not have
edges in between.
Parameters
----------
k : int
The desired number of negative samples to generate per edge.
exclude_self_loops : bool, optional
Whether to exclude self-loops from negative samples. (Default: True)
replace : bool, optional
Whether to sample with replacement. Setting it to True will make things
faster. (Default: True)
redundancy : float, optional
Indicates how much more negative samples to actually generate during rejection sampling
before finding the unique pairs.
Increasing it will increase the likelihood of getting :attr:`k` negative samples
per edge, but will also take more time and memory.
(Default: automatically determined by the density of graph)
Notes
-----
This negative sampler will try to generate as many negative samples as possible, but
it may rarely return less than :attr:`k` negative samples per edge.
This is more likely to happen if a graph is so small or dense that not many unique
negative samples exist.
Examples
--------
>>> g = dgl.graph(([0, 1, 2], [1, 2, 3]))
>>> neg_sampler = dgl.dataloading.negative_sampler.GlobalUniform(2, True)
>>> neg_sampler(g, torch.LongTensor([0, 1]))
(tensor([0, 1, 3, 2]), tensor([2, 0, 2, 1]))
"""
def __init__(self, k, exclude_self_loops=True, replace=False, redundancy=None):
self.k = k
self.exclude_self_loops = exclude_self_loops
self.replace = replace
self.redundancy = redundancy
def _generate(self, g, eids, canonical_etype):
return global_uniform_negative_sampling(
g, len(eids) * self.k, self.exclude_self_loops, self.replace,
canonical_etype, self.redundancy)
python/dgl/dataloading/pytorch/dataloader.py → python/dgl/_dataloading/pytorch/dataloader.py
View file @ 701b4fcc
......@@ -10,7 +10,6 @@ from torch.utils.data.distributed import DistributedSampler
import torch.distributed as dist
from ..dataloader import NodeCollator, EdgeCollator, GraphCollator, SubgraphIterator
from ...distributed import DistGraph
from ...distributed import DistDataLoader
from ...ndarray import NDArray as DGLNDArray
from ... import backend as F
from ...base import DGLError
......@@ -26,6 +25,10 @@ PYTORCH_VER = LooseVersion(th.__version__)
PYTORCH_16 = PYTORCH_VER >= LooseVersion("1.6.0")
PYTORCH_17 = PYTORCH_VER >= LooseVersion("1.7.0")
def _check_graph_type(g):
if isinstance(g, DistGraph):
raise TypeError("Please use DistNodeDataLoader or DistEdgeDataLoader for DistGraph")
def _create_dist_sampler(dataset, dataloader_kwargs, ddp_seed):
# Note: will change the content of dataloader_kwargs
dist_sampler_kwargs = {'shuffle': dataloader_kwargs['shuffle']}
......@@ -166,14 +169,6 @@ class _ScalarDataBatcher(th.utils.data.IterableDataset):
"""Set epoch number for distributed training."""
self.epoch = epoch
def _remove_kwargs_dist(kwargs):
if 'num_workers' in kwargs:
del kwargs['num_workers']
if 'pin_memory' in kwargs:
del kwargs['pin_memory']
print('Distributed DataLoader does not support pin_memory')
return kwargs
# The following code is a fix to the PyTorch-specific issue in
# https://github.com/dmlc/dgl/issues/2137
#
......@@ -290,14 +285,14 @@ def _restore_storages(subgs, g):
_restore_subgraph_storage(subg, g)
class _NodeCollator(NodeCollator):
def collate(self, items):
def collate(self, items): # pylint: disable=missing-docstring
# input_nodes, output_nodes, blocks
result = super().collate(items)
_pop_storages(result[-1], self.g)
return result
class _EdgeCollator(EdgeCollator):
def collate(self, items):
def collate(self, items): # pylint: disable=missing-docstring
if self.negative_sampler is None:
# input_nodes, pair_graph, blocks
result = super().collate(items)
......@@ -381,10 +376,10 @@ def _background_node_dataloader(dl_iter, g, device, results, load_input, load_ou
class _NodeDataLoaderIter:
def __init__(self, node_dataloader):
def __init__(self, node_dataloader, iter_):
self.device = node_dataloader.device
self.node_dataloader = node_dataloader
self.iter_ = iter(node_dataloader.dataloader)
self.iter_ = iter_
self.async_load = node_dataloader.async_load and (
F.device_type(self.device) == 'cuda')
if self.async_load:
......@@ -418,10 +413,10 @@ class _NodeDataLoaderIter:
return input_nodes, output_nodes, blocks
class _EdgeDataLoaderIter:
def __init__(self, edge_dataloader):
def __init__(self, edge_dataloader, iter_):
self.device = edge_dataloader.device
self.edge_dataloader = edge_dataloader
self.iter_ = iter(edge_dataloader.dataloader)
self.iter_ = iter_
# Make this an iterator for PyTorch Lightning compatibility
def __iter__(self):
......@@ -441,9 +436,9 @@ class _EdgeDataLoaderIter:
return result
class _GraphDataLoaderIter:
def __init__(self, graph_dataloader):
def __init__(self, graph_dataloader, iter_):
self.dataloader = graph_dataloader
self.iter_ = iter(graph_dataloader.dataloader)
self.iter_ = iter_
def __iter__(self):
return self
......@@ -490,14 +485,9 @@ def _init_dataloader(collator, device, dataloader_kwargs, use_ddp, ddp_seed):
else:
dist_sampler = None
dataloader = DataLoader(
dataset,
collate_fn=collator.collate,
**dataloader_kwargs)
return use_scalar_batcher, scalar_batcher, dataset, collator, dist_sampler
return use_scalar_batcher, scalar_batcher, dataloader, dist_sampler
class NodeDataLoader:
class NodeDataLoader(DataLoader):
"""PyTorch dataloader for batch-iterating over a set of nodes, generating the list
of message flow graphs (MFGs) as computation dependency of the said minibatch.
......@@ -600,6 +590,7 @@ class NodeDataLoader:
def __init__(self, g, nids, graph_sampler, device=None, use_ddp=False, ddp_seed=0,
load_input=None, load_output=None, async_load=False, **kwargs):
_check_graph_type(g)
collator_kwargs = {}
dataloader_kwargs = {}
for k, v in kwargs.items():
......@@ -608,65 +599,42 @@ class NodeDataLoader:
else:
dataloader_kwargs[k] = v
if isinstance(g, DistGraph):
if device is None:
# for the distributed case default to the CPU
device = 'cpu'
assert device == 'cpu', 'Only cpu is supported in the case of a DistGraph.'
# Distributed DataLoader currently does not support heterogeneous graphs
# and does not copy features. Fallback to normal solution
self.collator = NodeCollator(g, nids, graph_sampler, **collator_kwargs)
_remove_kwargs_dist(dataloader_kwargs)
self.dataloader = DistDataLoader(self.collator.dataset,
collate_fn=self.collator.collate,
**dataloader_kwargs)
self.is_distributed = True
else:
if device is None:
# default to the same device the graph is on
device = th.device(g.device)
if not g.is_homogeneous:
if load_input or load_output:
raise DGLError('load_input/load_output not supported for heterograph yet.')
self.load_input = {} if load_input is None else load_input
self.load_output = {} if load_output is None else load_output
self.async_load = async_load
# if the sampler supports it, tell it to output to the specified device.
# But if async_load is enabled, set_output_context should be skipped as
# we'd like to avoid any graph/data transfer graphs across devices in
# sampler. Such transfer will be handled in dataloader.
num_workers = dataloader_kwargs.get('num_workers', 0)
if ((not async_load) and
callable(getattr(graph_sampler, "set_output_context", None)) and
num_workers == 0):
graph_sampler.set_output_context(to_dgl_context(device))
self.collator = _NodeCollator(g, nids, graph_sampler, **collator_kwargs)
self.use_scalar_batcher, self.scalar_batcher, self.dataloader, self.dist_sampler = \
_init_dataloader(self.collator, device, dataloader_kwargs, use_ddp, ddp_seed)
if device is None:
# default to the same device the graph is on
device = th.device(g.device)
if not g.is_homogeneous:
if load_input or load_output:
raise DGLError('load_input/load_output not supported for heterograph yet.')
self.load_input = {} if load_input is None else load_input
self.load_output = {} if load_output is None else load_output
self.async_load = async_load
# if the sampler supports it, tell it to output to the specified device.
# But if async_load is enabled, set_output_context should be skipped as
# we'd like to avoid any graph/data transfer graphs across devices in
# sampler. Such transfer will be handled in dataloader.
num_workers = dataloader_kwargs.get('num_workers', 0)
if ((not async_load) and
callable(getattr(graph_sampler, "set_output_context", None)) and
num_workers == 0):
graph_sampler.set_output_context(to_dgl_context(device))
self.collator = _NodeCollator(g, nids, graph_sampler, **collator_kwargs)
self.use_scalar_batcher, self.scalar_batcher, self.dataloader, self.dist_sampler = \
_init_dataloader(self.collator, device, dataloader_kwargs, use_ddp, ddp_seed)
self.use_ddp = use_ddp
self.is_distributed = False
self.use_ddp = use_ddp
self.is_distributed = False
# Precompute the CSR and CSC representations so each subprocess does not
# duplicate.
if num_workers > 0:
g.create_formats_()
# Precompute the CSR and CSC representations so each subprocess does not
# duplicate.
if num_workers > 0:
g.create_formats_()
self.device = device
def __iter__(self):
"""Return the iterator of the data loader."""
if self.is_distributed:
# Directly use the iterator of DistDataLoader, which doesn't copy features anyway.
return iter(self.dataloader)
else:
return _NodeDataLoaderIter(self)
def __len__(self):
"""Return the number of batches of the data loader."""
return len(self.dataloader)
return _NodeDataLoaderIter(self, super().__iter__())
def set_epoch(self, epoch):
"""Sets the epoch number for the underlying sampler which ensures all replicas
......@@ -689,7 +657,7 @@ class NodeDataLoader:
else:
raise DGLError('set_epoch is only available when use_ddp is True.')
class EdgeDataLoader:
class EdgeDataLoader(DataLoader):
"""PyTorch dataloader for batch-iterating over a set of edges, generating the list
of message flow graphs (MFGs) as computation dependency of the said minibatch for
edge classification, edge regression, and link prediction.
......@@ -897,8 +865,9 @@ class EdgeDataLoader:
* Link prediction on heterogeneous graph: RGCN for link prediction.
"""
collator_arglist = inspect.getfullargspec(EdgeCollator).args
def __init__(self, g, eids, graph_sampler, device='cpu', use_ddp=False, ddp_seed=0, **kwargs):
def __init__(self, g, eids, graph_sampler, device='cpu', use_ddp=False, ddp_seed=0,
**kwargs):
_check_graph_type(g)
collator_kwargs = {}
dataloader_kwargs = {}
for k, v in kwargs.items():
......@@ -907,53 +876,30 @@ class EdgeDataLoader:
else:
dataloader_kwargs[k] = v
if isinstance(g, DistGraph):
if device is None:
# for the distributed case default to the CPU
device = 'cpu'
assert device == 'cpu', 'Only cpu is supported in the case of a DistGraph.'
# Distributed DataLoader currently does not support heterogeneous graphs
# and does not copy features. Fallback to normal solution
self.collator = EdgeCollator(g, eids, graph_sampler, **collator_kwargs)
_remove_kwargs_dist(dataloader_kwargs)
self.dataloader = DistDataLoader(self.collator.dataset,
collate_fn=self.collator.collate,
**dataloader_kwargs)
self.is_distributed = True
else:
if device is None:
# default to the same device the graph is on
device = th.device(g.device)
# if the sampler supports it, tell it to output to the
# specified device
num_workers = dataloader_kwargs.get('num_workers', 0)
if callable(getattr(graph_sampler, "set_output_context", None)) and num_workers == 0:
graph_sampler.set_output_context(to_dgl_context(device))
# if the sampler supports it, tell it to output to the
# specified device
num_workers = dataloader_kwargs.get('num_workers', 0)
if callable(getattr(graph_sampler, "set_output_context", None)) and num_workers == 0:
graph_sampler.set_output_context(to_dgl_context(device))
self.collator = _EdgeCollator(g, eids, graph_sampler, **collator_kwargs)
self.use_scalar_batcher, self.scalar_batcher, self.dataloader, self.dist_sampler = \
_init_dataloader(self.collator, device, dataloader_kwargs, use_ddp, ddp_seed)
self.use_ddp = use_ddp
self.is_distributed = False
self.collator = EdgeCollator(g, eids, graph_sampler, **collator_kwargs)
self.use_scalar_batcher, self.scalar_batcher, dataset, collator, self.dist_sampler = \
_init_dataloader(self.collator, device, dataloader_kwargs, use_ddp, ddp_seed)
self.use_ddp = use_ddp
super().__init__(dataset, collate_fn=collator.collate, **dataloader_kwargs)
# Precompute the CSR and CSC representations so each subprocess does not duplicate.
if num_workers > 0:
g.create_formats_()
# Precompute the CSR and CSC representations so each subprocess does not duplicate.
if num_workers > 0:
g.create_formats_()
self.device = device
def __iter__(self):
"""Return the iterator of the data loader."""
if self.is_distributed:
# Directly use the iterator of DistDataLoader, which doesn't copy features anyway.
return iter(self.dataloader)
else:
return _EdgeDataLoaderIter(self)
def __len__(self):
"""Return the number of batches of the data loader."""
return len(self.dataloader)
return _EdgeDataLoaderIter(self, super().__iter__())
def set_epoch(self, epoch):
"""Sets the epoch number for the underlying sampler which ensures all replicas
......@@ -976,7 +922,7 @@ class EdgeDataLoader:
else:
raise DGLError('set_epoch is only available when use_ddp is True.')
class GraphDataLoader:
class GraphDataLoader(DataLoader):
"""PyTorch dataloader for batch-iterating over a set of graphs, generating the batched
graph and corresponding label tensor (if provided) of the said minibatch.
......@@ -1023,7 +969,6 @@ class GraphDataLoader:
... train_on(batched_graph, labels)
"""
collator_arglist = inspect.getfullargspec(GraphCollator).args
def __init__(self, dataset, collate_fn=None, use_ddp=False, ddp_seed=0, **kwargs):
collator_kwargs = {}
dataloader_kwargs = {}
......@@ -1058,14 +1003,11 @@ class GraphDataLoader:
if use_ddp:
self.dist_sampler = _create_dist_sampler(dataset, dataloader_kwargs, ddp_seed)
dataloader_kwargs['sampler'] = self.dist_sampler
self.dataloader = DataLoader(dataset=dataset,
collate_fn=self.collate,
**dataloader_kwargs)
super().__init__(dataset, collate_fn=self.collate, **dataloader_kwargs)
def __iter__(self):
"""Return the iterator of the data loader."""
return _GraphDataLoaderIter(self)
return _GraphDataLoaderIter(self, super().__iter__())
def __len__(self):
"""Return the number of batches of the data loader."""
......
python/dgl/_dataloading/shadow.py 0 → 100644
View file @ 701b4fcc
"""ShaDow-GNN subgraph samplers."""
from ..utils import prepare_tensor_or_dict
from ..base import NID
from .. import transform
from ..sampling import sample_neighbors
from .neighbor import NeighborSamplingMixin
from .dataloader import exclude_edges, Sampler
class ShaDowKHopSampler(NeighborSamplingMixin, Sampler):
"""K-hop subgraph sampler used by
`ShaDow-GNN <https://arxiv.org/abs/2012.01380>`__.
It performs node-wise neighbor sampling but instead of returning a list of
MFGs, it returns a single subgraph induced by all the sampled nodes. The
seed nodes from which the neighbors are sampled will appear the first in the
induced nodes of the subgraph.
This is used in conjunction with :class:`dgl.dataloading.pytorch.NodeDataLoader`
and :class:`dgl.dataloading.pytorch.EdgeDataLoader`.
Parameters
----------
fanouts : list[int] or list[dict[etype, int]]
List of neighbors to sample per edge type for each GNN layer, with the i-th
element being the fanout for the i-th GNN layer.
If only a single integer is provided, DGL assumes that every edge type
will have the same fanout.
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
prob : str, optional
If given, the probability of each neighbor being sampled is proportional
to the edge feature value with the given name in ``g.edata``. The feature must be
a scalar on each edge.
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 5, 10, 15 neighbors for
the first, second, and third layer respectively (assuming the backend is PyTorch):
>>> g = dgl.data.CoraFullDataset()[0]
>>> sampler = dgl.dataloading.ShaDowKHopSampler([5, 10, 15])
>>> dataloader = dgl.dataloading.NodeDataLoader(
... g, torch.arange(g.num_nodes()), sampler,
... batch_size=5, shuffle=True, drop_last=False, num_workers=4)
>>> for input_nodes, output_nodes, (subgraph,) in dataloader:
... print(subgraph)
... assert torch.equal(input_nodes, subgraph.ndata[dgl.NID])
... assert torch.equal(input_nodes[:output_nodes.shape[0]], output_nodes)
... break
Graph(num_nodes=529, num_edges=3796,
ndata_schemes={'label': Scheme(shape=(), dtype=torch.int64),
'feat': Scheme(shape=(8710,), dtype=torch.float32),
'_ID': Scheme(shape=(), dtype=torch.int64)}
edata_schemes={'_ID': Scheme(shape=(), dtype=torch.int64)})
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.dataloading.ShaDowKHopSampler([
... {('user', 'follows', 'user'): 5,
... ('user', 'plays', 'game'): 4,
... ('game', 'played-by', 'user'): 3}] * 3)
If you would like non-uniform neighbor sampling:
>>> g.edata['p'] = torch.rand(g.num_edges()) # any non-negative 1D vector works
>>> sampler = dgl.dataloading.MultiLayerNeighborSampler([5, 10, 15], prob='p')
"""
def __init__(self, fanouts, replace=False, prob=None, output_ctx=None):
super().__init__(output_ctx)
self.fanouts = fanouts
self.replace = replace
self.prob = prob
self.set_output_context(output_ctx)
def sample(self, g, seed_nodes, exclude_eids=None):
self._build_fanout(len(self.fanouts), g)
self._build_prob_arrays(g)
seed_nodes = prepare_tensor_or_dict(g, seed_nodes, 'seed nodes')
output_nodes = seed_nodes
for i in range(len(self.fanouts)):
fanout = self.fanouts[i]
frontier = sample_neighbors(
g, seed_nodes, fanout, replace=self.replace, prob=self.prob_arrays)
block = transform.to_block(frontier, seed_nodes)
seed_nodes = block.srcdata[NID]
subg = g.subgraph(seed_nodes, relabel_nodes=True)
subg = exclude_edges(subg, exclude_eids, self.output_device)
return seed_nodes, output_nodes, [subg]
python/dgl/dataloading/__init__.py
View file @ 701b4fcc
......@@ -14,15 +14,12 @@ Read the user guide :ref:`guide-minibatch`.
This package is experimental and the interfaces may be subject
to changes in future releases. It currently only has implementations in PyTorch.
"""
from .neighbor import *
from .dataloader import *
from .. import backend as F
from .neighbor_sampler import *
from .cluster_gcn import *
from .shadow import *
from .base import *
from . import negative_sampler
from .async_transferer import AsyncTransferer
from .. import backend as F
if F.get_preferred_backend() == 'pytorch':
from .pytorch import *
from .dataloader import *
from .dist_dataloader import *
python/dgl/dataloading/base.py 0 → 100644
View file @ 701b4fcc
"""Base classes and functionalities for dataloaders"""
from collections import Mapping
from ..base import NID, EID
from ..convert import heterograph
from .. import backend as F
from ..transform import compact_graphs
from ..frame import LazyFeature
from ..utils import recursive_apply
def _set_lazy_features(x, xdata, feature_names):
if feature_names is None:
return
if not isinstance(feature_names, Mapping):
xdata.update({k: LazyFeature(k) for k in feature_names})
else:
for type_, names in feature_names.items():
x[type_].data.update({k: LazyFeature(k) for k in names})
def set_node_lazy_features(g, feature_names):
"""Set lazy features for ``g.ndata`` if :attr:`feature_names` is a list of strings,
or ``g.nodes[ntype].data`` if :attr:`feature_names` is a dict of list of strings.
"""
return _set_lazy_features(g.nodes, g.ndata, feature_names)
def set_edge_lazy_features(g, feature_names):
"""Set lazy features for ``g.edata`` if :attr:`feature_names` is a list of strings,
or ``g.edges[etype].data`` if :attr:`feature_names` is a dict of list of strings.
"""
return _set_lazy_features(g.edges, g.edata, feature_names)
def set_src_lazy_features(g, feature_names):
"""Set lazy features for ``g.srcdata`` if :attr:`feature_names` is a list of strings,
or ``g.srcnodes[srctype].data`` if :attr:`feature_names` is a dict of list of strings.
"""
return _set_lazy_features(g.srcnodes, g.srcdata, feature_names)
def set_dst_lazy_features(g, feature_names):
"""Set lazy features for ``g.dstdata`` if :attr:`feature_names` is a list of strings,
or ``g.dstnodes[dsttype].data`` if :attr:`feature_names` is a dict of list of strings.
"""
return _set_lazy_features(g.dstnodes, g.dstdata, feature_names)
class BlockSampler(object):
"""BlockSampler is an abstract class assuming to take in a set of nodes whose
outputs are to compute, and return a list of blocks.
Moreover, it assumes that the input node features will be put in the first block's
``srcdata``, the output node labels will be put in the last block's ``dstdata``, and
the edge data will be put in all the blocks' ``edata``.
"""
def __init__(self, prefetch_node_feats=None, prefetch_labels=None,
prefetch_edge_feats=None, output_device=None):
self.prefetch_node_feats = prefetch_node_feats or []
self.prefetch_labels = prefetch_labels or []
self.prefetch_edge_feats = prefetch_edge_feats or []
self.output_device = output_device
def sample_blocks(self, g, seed_nodes, exclude_eids=None):
"""Generates a list of blocks from the given seed nodes.
This function must return a triplet where the first element is the input node IDs
for the first GNN layer (a tensor or a dict of tensors for heterogeneous graphs),
the second element is the output node IDs for the last GNN layer, and the third
element is the said list of blocks.
"""
raise NotImplementedError
def assign_lazy_features(self, result):
"""Assign lazy features for prefetching."""
# A LazyFeature is a placeholder telling the dataloader where and which IDs
# to prefetch. It has the signature LazyFeature(name, id_). id_ can be None
# if the LazyFeature is set into one of the subgraph's ``xdata``, in which case the
# dataloader will infer from the subgraph's ``xdata[dgl.NID]`` (or ``xdata[dgl.EID]``
# if the LazyFeature is set as edge features).
#
# If you want to prefetch things other than ndata and edata, you can also
# return a LazyFeature(name, id_). If a LazyFeature is returned in places other than
# in a graph's ndata/edata/srcdata/dstdata, the DataLoader will prefetch it
# from its dictionary ``other_data``.
# For instance, you can run
#
# return blocks, LazyFeature('other_feat', id_)
#
# To make it work with the sampler returning the stuff above, your dataloader
# needs to have the following
#
# dataloader.attach_data('other_feat', tensor)
#
# Then you can run
#
# for blocks, other_feat in dataloader:
# train_on(blocks, other_feat)
input_nodes, output_nodes, blocks = result
set_src_lazy_features(blocks[0], self.prefetch_node_feats)
set_dst_lazy_features(blocks[-1], self.prefetch_labels)
for block in blocks:
set_edge_lazy_features(block, self.prefetch_edge_feats)
return input_nodes, output_nodes, blocks
def sample(self, g, seed_nodes):
"""Sample a list of blocks from the given seed nodes."""
result = self.sample_blocks(g, seed_nodes)
return self.assign_lazy_features(result)
def _find_exclude_eids_with_reverse_id(g, eids, reverse_eid_map):
if isinstance(eids, Mapping):
eids = {g.to_canonical_etype(k): v for k, v in eids.items()}
exclude_eids = {
k: F.cat([v, F.gather_row(reverse_eid_map[k], v)], 0)
for k, v in eids.items()}
else:
exclude_eids = F.cat([eids, F.gather_row(reverse_eid_map, eids)], 0)
return exclude_eids
def _find_exclude_eids_with_reverse_types(g, eids, reverse_etype_map):
exclude_eids = {g.to_canonical_etype(k): v for k, v in eids.items()}
reverse_etype_map = {
g.to_canonical_etype(k): g.to_canonical_etype(v)
for k, v in reverse_etype_map.items()}
exclude_eids.update({reverse_etype_map[k]: v for k, v in exclude_eids.items()})
return exclude_eids
def _find_exclude_eids(g, exclude_mode, eids, **kwargs):
if exclude_mode is None:
return None
elif F.is_tensor(exclude_mode) or (
isinstance(exclude_mode, Mapping) and
all(F.is_tensor(v) for v in exclude_mode.values())):
return exclude_mode
elif exclude_mode == 'self':
return eids
elif exclude_mode == 'reverse_id':
return _find_exclude_eids_with_reverse_id(g, eids, kwargs['reverse_eid_map'])
elif exclude_mode == 'reverse_types':
return _find_exclude_eids_with_reverse_types(g, eids, kwargs['reverse_etype_map'])
else:
raise ValueError('unsupported mode {}'.format(exclude_mode))
def find_exclude_eids(g, seed_edges, exclude, reverse_eids=None, reverse_etypes=None,
output_device=None):
"""Find all edge IDs to exclude according to :attr:`exclude_mode`.
Parameters
----------
g : DGLGraph
The graph.
exclude_mode : str, optional
Can be either of the following,
None (default)
Does not exclude any edge.
Tensor or dict[etype, Tensor]
Exclude the given edge IDs.
'self'
Exclude the given edges themselves but nothing else.
'reverse_id'
Exclude all edges specified in ``eids``, as well as their reverse edges
of the same edge type.
The mapping from each edge ID to its reverse edge ID is specified in
the keyword argument ``reverse_eid_map``.
This mode assumes that the reverse of an edge with ID ``e`` and type
``etype`` will have ID ``reverse_eid_map[e]`` and type ``etype``.
'reverse_types'
Exclude all edges specified in ``eids``, as well as their reverse
edges of the corresponding edge types.
The mapping from each edge type to its reverse edge type is specified
in the keyword argument ``reverse_etype_map``.
This mode assumes that the reverse of an edge with ID ``e`` and type ``etype``
will have ID ``e`` and type ``reverse_etype_map[etype]``.
eids : Tensor or dict[etype, Tensor]
The edge IDs.
reverse_eids : Tensor or dict[etype, Tensor]
The mapping from edge ID to its reverse edge ID.
reverse_etypes : dict[etype, etype]
The mapping from edge etype to its reverse edge type.
output_device : device
The device of the output edge IDs.
"""
exclude_eids = _find_exclude_eids(
g,
exclude,
seed_edges,
reverse_eid_map=reverse_eids,
reverse_etype_map=reverse_etypes)
if exclude_eids is not None:
exclude_eids = recursive_apply(
exclude_eids, lambda x: x.to(output_device))
return exclude_eids
class EdgeBlockSampler(object):
"""Adapts a :class:`BlockSampler` object's :attr:`sample` method for edge
classification and link prediction.
"""
def __init__(self, block_sampler, exclude=None, reverse_eids=None,
reverse_etypes=None, negative_sampler=None, prefetch_node_feats=None,
prefetch_labels=None, prefetch_edge_feats=None):
self.reverse_eids = reverse_eids
self.reverse_etypes = reverse_etypes
self.exclude = exclude
self.block_sampler = block_sampler
self.negative_sampler = negative_sampler
self.prefetch_node_feats = prefetch_node_feats or []
self.prefetch_labels = prefetch_labels or []
self.prefetch_edge_feats = prefetch_edge_feats or []
self.output_device = block_sampler.output_device
def _build_neg_graph(self, g, seed_edges):
neg_srcdst = self.negative_sampler(g, seed_edges)
if not isinstance(neg_srcdst, Mapping):
assert len(g.canonical_etypes) == 1, \
'graph has multiple or no edge types; '\
'please return a dict in negative sampler.'
neg_srcdst = {g.canonical_etypes[0]: neg_srcdst}
dtype = F.dtype(list(neg_srcdst.values())[0][0])
neg_edges = {
etype: neg_srcdst.get(etype, (F.tensor([], dtype), F.tensor([], dtype)))
for etype in g.canonical_etypes}
neg_pair_graph = heterograph(
neg_edges, {ntype: g.num_nodes(ntype) for ntype in g.ntypes})
return neg_pair_graph
def assign_lazy_features(self, result):
"""Assign lazy features for prefetching."""
pair_graph = result[1]
blocks = result[-1]
set_src_lazy_features(blocks[0], self.prefetch_node_feats)
set_edge_lazy_features(pair_graph, self.prefetch_labels)
for block in blocks:
set_edge_lazy_features(block, self.prefetch_edge_feats)
# In-place updates
return result
def sample(self, g, seed_edges):
"""Samples a list of blocks, as well as a subgraph containing the sampled
edges from the original graph.
If :attr:`negative_sampler` is given, also returns another graph containing the
negative pairs as edges.
"""
exclude = self.exclude
pair_graph = g.edge_subgraph(
seed_edges, relabel_nodes=False, output_device=self.output_device)
eids = pair_graph.edata[EID]
if self.negative_sampler is not None:
neg_graph = self._build_neg_graph(g, seed_edges)
pair_graph, neg_graph = compact_graphs([pair_graph, neg_graph])
else:
pair_graph = compact_graphs(pair_graph)
pair_graph.edata[EID] = eids
seed_nodes = pair_graph.ndata[NID]
exclude_eids = find_exclude_eids(
g, seed_edges, exclude, self.reverse_eids, self.reverse_etypes,
self.output_device)
input_nodes, _, blocks = self.block_sampler.sample_blocks(g, seed_nodes, exclude_eids)
if self.negative_sampler is None:
return self.assign_lazy_features((input_nodes, pair_graph, blocks))
else:
return self.assign_lazy_features((input_nodes, pair_graph, neg_graph, blocks))
python/dgl/dataloading/cluster_gcn.py
View file @ 701b4fcc
"""Cluster-GCN subgraph iterators."""
"""Cluster-GCN samplers."""
import os
import pickle
import numpy as np
from ..transform import metis_partition_assignment
from .. import backend as F
from .dataloader import SubgraphIterator
from ..base import DGLError
from ..partition import metis_partition_assignment
from .base import set_node_lazy_features, set_edge_lazy_features
class ClusterGCNSubgraphIterator(SubgraphIterator):
"""Subgraph sampler following that of ClusterGCN.
class ClusterGCNSampler(object):
"""Cluster-GCN sampler.
This sampler first partitions the graph with METIS partitioning, then it caches the nodes of
each partition to a file within the given cache directory.
This is used in conjunction with :class:`dgl.dataloading.pytorch.GraphDataLoader`.
This is used in conjunction with :class:`dgl.dataloading.DataLoader`.
Notes
-----
......@@ -23,61 +24,53 @@ class ClusterGCNSubgraphIterator(SubgraphIterator):
----------
g : DGLGraph
The original graph.
num_partitions : int
k : int
The number of partitions.
cache_directory : str
cache_path : str
The path to the cache directory for storing the partition result.
refresh : bool
If True, recompute the partition.
Examples
--------
Assuming that you have a graph ``g``:
>>> sgiter = dgl.dataloading.ClusterGCNSubgraphIterator(
... g, num_partitions=100, cache_directory='.', refresh=True)
>>> dataloader = dgl.dataloading.GraphDataLoader(sgiter, batch_size=4, num_workers=0)
>>> for subgraph_batch in dataloader:
... train_on(subgraph_batch)
"""
def __init__(self, g, num_partitions, cache_directory, refresh=False):
if os.name == 'nt':
raise NotImplementedError("METIS partitioning is not supported on Windows yet.")
super().__init__(g)
# First see if the cache is already there. If so, directly read from cache.
if not refresh and self._load_parts(cache_directory):
return
# Otherwise, build the cache.
assignment = F.asnumpy(metis_partition_assignment(g, num_partitions))
self._save_parts(assignment, cache_directory)
def _cache_file_path(self, cache_directory):
return os.path.join(cache_directory, 'cluster_gcn_cache')
def _load_parts(self, cache_directory):
path = self._cache_file_path(cache_directory)
if not os.path.exists(path):
return False
with open(path, 'rb') as file_:
self.part_indptr, self.part_indices = pickle.load(file_)
return True
def _save_parts(self, assignment, cache_directory):
os.makedirs(cache_directory, exist_ok=True)
self.part_indices = np.argsort(assignment)
num_nodes_per_part = np.bincount(assignment)
self.part_indptr = np.insert(np.cumsum(num_nodes_per_part), 0, 0)
with open(self._cache_file_path(cache_directory), 'wb') as file_:
pickle.dump((self.part_indptr, self.part_indices), file_)
def __len__(self):
return self.part_indptr.shape[0] - 1
def __getitem__(self, i):
nodes = self.part_indices[self.part_indptr[i]:self.part_indptr[i+1]]
return self.g.subgraph(nodes)
def __init__(self, g, k, balance_ntypes=None, balance_edges=False, mode='k-way',
prefetch_node_feats=None, prefetch_edge_feats=None, output_device=None,
cache_path='cluster_gcn.pkl'):
if os.path.exists(cache_path):
try:
with open(cache_path, 'rb') as f:
self.partition_offset, self.partition_node_ids = pickle.load(f)
except (EOFError, TypeError, ValueError):
raise DGLError(
f'The contents in the cache file {cache_path} is invalid. '
f'Please remove the cache file {cache_path} or specify another path.')
if len(self.partition_offset) != k + 1:
raise DGLError(
f'Number of partitions in the cache does not match the value of k. '
f'Please remove the cache file {cache_path} or specify another path.')
if len(self.partition_node_ids) != g.num_nodes():
raise DGLError(
f'Number of nodes in the cache does not match the given graph. '
f'Please remove the cache file {cache_path} or specify another path.')
else:
partition_ids = metis_partition_assignment(
g, k, balance_ntypes=balance_ntypes, balance_edges=balance_edges, mode=mode)
partition_ids = F.asnumpy(partition_ids)
partition_node_ids = np.argsort(partition_ids)
partition_size = F.zerocopy_from_numpy(np.bincount(partition_ids, minlength=k))
partition_offset = F.zerocopy_from_numpy(np.insert(np.cumsum(partition_size), 0, 0))
partition_node_ids = F.zerocopy_from_numpy(partition_ids)
with open(cache_path, 'wb') as f:
pickle.dump((partition_offset, partition_node_ids), f)
self.partition_offset = partition_offset
self.partition_node_ids = partition_node_ids
self.prefetch_node_feats = prefetch_node_feats or []
self.prefetch_edge_feats = prefetch_edge_feats or []
self.output_device = output_device
def sample(self, g, partition_ids):
"""Samples a subgraph given a list of partition IDs."""
node_ids = F.cat([
self.partition_node_ids[self.partition_offset[i]:self.partition_offset[i+1]]
for i in F.asnumpy(partition_ids)], 0)
sg = g.subgraph(node_ids, relabel_nodes=True, output_device=self.output_device)
set_node_lazy_features(sg, self.prefetch_node_feats)
set_edge_lazy_features(sg, self.prefetch_edge_feats)
return sg
python/dgl/dataloading/dataloader.py
View file @ 701b4fcc
"""Data loaders"""
"""DGL PyTorch DataLoaders"""
from collections.abc import Mapping, Sequence
from abc import ABC, abstractproperty, abstractmethod
from queue import Queue
import itertools
import threading
from distutils.version import LooseVersion
import random
import math
import inspect
import re
import numpy as np
from .. import transform
from ..base import NID, EID
from .. import backend as F
from .. import utils
from ..batch import batch
from ..convert import heterograph
from ..heterograph import DGLHeteroGraph as DGLGraph
from ..distributed.dist_graph import DistGraph
from ..utils import to_device
def _tensor_or_dict_to_numpy(ids):
if isinstance(ids, Mapping):
return {k: F.zerocopy_to_numpy(v) for k, v in ids.items()}
else:
return F.zerocopy_to_numpy(ids)
def _locate_eids_to_exclude(frontier_parent_eids, exclude_eids):
"""Find the edges whose IDs in parent graph appeared in exclude_eids.
import torch
import torch.distributed as dist
from torch.utils.data.distributed import DistributedSampler
from ..base import NID, EID, dgl_warning
from ..batch import batch as batch_graphs
from ..heterograph import DGLHeteroGraph
from .. import ndarray as nd
from ..utils import (
recursive_apply, ExceptionWrapper, recursive_apply_pair, set_num_threads,
create_shared_mem_array, get_shared_mem_array)
from ..frame import LazyFeature
from ..storages import wrap_storage
from .base import BlockSampler, EdgeBlockSampler
from .. import backend as F
Note that both arguments are numpy arrays or numpy dicts.
class _TensorizedDatasetIter(object):
def __init__(self, dataset, batch_size, drop_last, mapping_keys):
self.dataset = dataset
self.batch_size = batch_size
self.drop_last = drop_last
self.mapping_keys = mapping_keys
self.index = 0
# For PyTorch Lightning compatibility
def __iter__(self):
return self
def _next_indices(self):
num_items = self.dataset.shape[0]
if self.index >= num_items:
raise StopIteration
end_idx = self.index + self.batch_size
if end_idx > num_items:
if self.drop_last:
raise StopIteration
end_idx = num_items
batch = self.dataset[self.index:end_idx]
self.index += self.batch_size
return batch
def __next__(self):
batch = self._next_indices()
if self.mapping_keys is None:
return batch
# convert the type-ID pairs to dictionary
type_ids = batch[:, 0]
indices = batch[:, 1]
type_ids_sortidx = torch.argsort(type_ids)
type_ids = type_ids[type_ids_sortidx]
indices = indices[type_ids_sortidx]
type_id_uniq, type_id_count = torch.unique_consecutive(type_ids, return_counts=True)
type_id_uniq = type_id_uniq.tolist()
type_id_offset = type_id_count.cumsum(0).tolist()
type_id_offset.insert(0, 0)
id_dict = {
self.mapping_keys[type_id_uniq[i]]: indices[type_id_offset[i]:type_id_offset[i+1]]
for i in range(len(type_id_uniq))}
return id_dict
def _get_id_tensor_from_mapping(indices, device, keys):
lengths = torch.LongTensor([
(indices[k].shape[0] if k in indices else 0) for k in keys], device=device)
type_ids = torch.arange(len(keys), device=device).repeat_interleave(lengths)
all_indices = torch.cat([indices[k] for k in keys if k in indices])
return torch.stack([type_ids, all_indices], 1)
def _divide_by_worker(dataset):
num_samples = dataset.shape[0]
worker_info = torch.utils.data.get_worker_info()
if worker_info:
chunk_size = num_samples // worker_info.num_workers
left_over = num_samples % worker_info.num_workers
start = (chunk_size * worker_info.id) + min(left_over, worker_info.id)
end = start + chunk_size + (worker_info.id < left_over)
assert worker_info.id < worker_info.num_workers - 1 or end == num_samples
dataset = dataset[start:end]
return dataset
class TensorizedDataset(torch.utils.data.IterableDataset):
"""Custom Dataset wrapper that returns a minibatch as tensors or dicts of tensors.
When the dataset is on the GPU, this significantly reduces the overhead.
"""
if isinstance(frontier_parent_eids, Mapping):
result = {
k: np.isin(frontier_parent_eids[k], exclude_eids[k]).nonzero()[0]
for k in frontier_parent_eids.keys() if k in exclude_eids.keys()}
return {k: F.zerocopy_from_numpy(v) for k, v in result.items()}
else:
result = np.isin(frontier_parent_eids, exclude_eids).nonzero()[0]
return F.zerocopy_from_numpy(result)
class _EidExcluder():
def __init__(self, exclude_eids):
device = None
if isinstance(exclude_eids, Mapping):
for _, v in exclude_eids.items():
if device is None:
device = F.context(v)
break
def __init__(self, indices, batch_size, drop_last):
if isinstance(indices, Mapping):
self._mapping_keys = list(indices.keys())
self._device = next(iter(indices.values())).device
self._tensor_dataset = _get_id_tensor_from_mapping(
indices, self._device, self._mapping_keys)
else:
device = F.context(exclude_eids)
self._exclude_eids = None
self._filter = None
if device == F.cpu():
# TODO(nv-dlasalle): Once Filter is implemented for the CPU, we
# should just use that irregardless of the device.
self._exclude_eids = (
_tensor_or_dict_to_numpy(exclude_eids) if exclude_eids is not None else None)
self._tensor_dataset = indices
self._device = indices.device
self._mapping_keys = None
self.batch_size = batch_size
self.drop_last = drop_last
def shuffle(self):
"""Shuffle the dataset."""
# TODO: may need an in-place shuffle kernel
perm = torch.randperm(self._tensor_dataset.shape[0], device=self._device)
self._tensor_dataset[:] = self._tensor_dataset[perm]
def __iter__(self):
dataset = _divide_by_worker(self._tensor_dataset)
return _TensorizedDatasetIter(
dataset, self.batch_size, self.drop_last, self._mapping_keys)
def __len__(self):
num_samples = self._tensor_dataset.shape[0]
return (num_samples + (0 if self.drop_last else (self.batch_size - 1))) // self.batch_size
def _get_shared_mem_name(id_):
return f'ddp_{id_}'
def _generate_shared_mem_name_id():
for _ in range(3): # 3 trials
id_ = random.getrandbits(32)
name = _get_shared_mem_name(id_)
if not nd.exist_shared_mem_array(name):
return name, id_
raise DGLError('Unable to generate a shared memory array')
class DDPTensorizedDataset(torch.utils.data.IterableDataset):
"""Custom Dataset wrapper that returns a minibatch as tensors or dicts of tensors.
When the dataset is on the GPU, this significantly reduces the overhead.
This class additionally saves the index tensor in shared memory and therefore
avoids duplicating the same index tensor during shuffling.
"""
def __init__(self, indices, batch_size, drop_last, ddp_seed):
if isinstance(indices, Mapping):
self._mapping_keys = list(indices.keys())
else:
if isinstance(exclude_eids, Mapping):
self._filter = {k: utils.Filter(v) for k, v in exclude_eids.items()}
else:
self._filter = utils.Filter(exclude_eids)
self._mapping_keys = None
def _find_indices(self, parent_eids):
""" Find the set of edge indices to remove.
"""
if self._exclude_eids is not None:
parent_eids_np = _tensor_or_dict_to_numpy(parent_eids)
return _locate_eids_to_exclude(parent_eids_np, self._exclude_eids)
self.rank = dist.get_rank()
self.num_replicas = dist.get_world_size()
self.seed = ddp_seed
self.epoch = 0
self.batch_size = batch_size
self.drop_last = drop_last
if self.drop_last and len(indices) % self.num_replicas != 0:
self.num_samples = math.ceil((len(indices) - self.num_replicas) / self.num_replicas)
else:
assert self._filter is not None
if isinstance(parent_eids, Mapping):
located_eids = {k: self._filter[k].find_included_indices(parent_eids[k])
for k, v in parent_eids.items() if k in self._filter}
self.num_samples = math.ceil(len(indices) / self.num_replicas)
self.total_size = self.num_samples * self.num_replicas
# If drop_last is True, we create a shared memory array larger than the number
# of indices since we will need to pad it after shuffling to make it evenly
# divisible before every epoch. If drop_last is False, we create an array
# with the same size as the indices so we can trim it later.
self.shared_mem_size = self.total_size if not self.drop_last else len(indices)
self.num_indices = len(indices)
if self.rank == 0:
name, id_ = _generate_shared_mem_name_id()
if isinstance(indices, Mapping):
device = next(iter(indices.values())).device
id_tensor = _get_id_tensor_from_mapping(indices, device, self._mapping_keys)
self._tensor_dataset = create_shared_mem_array(
name, (self.shared_mem_size, 2), torch.int64)
self._tensor_dataset[:id_tensor.shape[0], :] = id_tensor
else:
located_eids = self._filter.find_included_indices(parent_eids)
return located_eids
def __call__(self, frontier):
parent_eids = frontier.edata[EID]
located_eids = self._find_indices(parent_eids)
if not isinstance(located_eids, Mapping):
# (BarclayII) If frontier already has a EID field and located_eids is empty,
# the returned graph will keep EID intact. Otherwise, EID will change
# to the mapping from the new graph to the old frontier.
# So we need to test if located_eids is empty, and do the remapping ourselves.
if len(located_eids) > 0:
frontier = transform.remove_edges(
frontier, located_eids, store_ids=True)
frontier.edata[EID] = F.gather_row(parent_eids, frontier.edata[EID])
else:
# (BarclayII) remove_edges only accepts removing one type of edges,
# so I need to keep track of the edge IDs left one by one.
new_eids = parent_eids.copy()
for k, v in located_eids.items():
if len(v) > 0:
frontier = transform.remove_edges(
frontier, v, etype=k, store_ids=True)
new_eids[k] = F.gather_row(parent_eids[k], frontier.edges[k].data[EID])
frontier.edata[EID] = new_eids
return frontier
def exclude_edges(subg, exclude_eids, device):
"""Find and remove from the subgraph the edges whose IDs in the parent
graph are given.
Parameters
----------
subg : DGLGraph
The subgraph. Must have ``dgl.EID`` field containing the original
edge IDs in the parent graph.
exclude_eids : Tensor or dict
The edge IDs to exclude.
device : device
The output device of the graph.
Returns
-------
DGLGraph
The new subgraph with edges removed. The ``dgl.EID`` field contains
the original edge IDs in the same parent graph.
"""
if exclude_eids is None:
return subg
if device is not None:
if isinstance(exclude_eids, Mapping):
exclude_eids = {k: F.copy_to(v, device) \
for k, v in exclude_eids.items()}
self._tensor_dataset = create_shared_mem_array(
name, (self.shared_mem_size,), torch.int64)
self._tensor_dataset[:len(indices)] = indices
self._device = self._tensor_dataset.device
meta_info = torch.LongTensor([id_, self._tensor_dataset.shape[0]])
else:
exclude_eids = F.copy_to(exclude_eids, device)
excluder = _EidExcluder(exclude_eids)
return subg if excluder is None else excluder(subg)
def _find_exclude_eids_with_reverse_id(g, eids, reverse_eid_map):
if isinstance(eids, Mapping):
eids = {g.to_canonical_etype(k): v for k, v in eids.items()}
exclude_eids = {
k: F.cat([v, F.gather_row(reverse_eid_map[k], v)], 0)
for k, v in eids.items()}
meta_info = torch.LongTensor([0, 0])
if dist.get_backend() == 'nccl':
# Use default CUDA device; PyTorch DDP required the users to set the CUDA
# device for each process themselves so calling .cuda() should be safe.
meta_info = meta_info.cuda()
dist.broadcast(meta_info, src=0)
if self.rank != 0:
id_, num_samples = meta_info.tolist()
name = _get_shared_mem_name(id_)
if isinstance(indices, Mapping):
indices_shared = get_shared_mem_array(name, (num_samples, 2), torch.int64)
else:
indices_shared = get_shared_mem_array(name, (num_samples,), torch.int64)
self._tensor_dataset = indices_shared
self._device = indices_shared.device
def shuffle(self):
"""Shuffles the dataset."""
# Only rank 0 does the actual shuffling. The other ranks wait for it.
if self.rank == 0:
self._tensor_dataset[:self.num_indices] = self._tensor_dataset[
torch.randperm(self.num_indices, device=self._device)]
if not self.drop_last:
# pad extra
self._tensor_dataset[self.num_indices:] = \
self._tensor_dataset[:self.total_size - self.num_indices]
dist.barrier()
def __iter__(self):
start = self.num_samples * self.rank
end = self.num_samples * (self.rank + 1)
dataset = _divide_by_worker(self._tensor_dataset[start:end])
return _TensorizedDatasetIter(
dataset, self.batch_size, self.drop_last, self._mapping_keys)
def __len__(self):
return (self.num_samples + (0 if self.drop_last else (self.batch_size - 1))) // \
self.batch_size
def _prefetch_update_feats(feats, frames, types, get_storage_func, id_name, device, pin_memory):
for tid, frame in enumerate(frames):
type_ = types[tid]
default_id = frame.get(id_name, None)
for key in frame.keys():
column = frame[key]
if isinstance(column, LazyFeature):
parent_key = column.name or key
if column.id_ is None and default_id is None:
raise DGLError(
'Found a LazyFeature with no ID specified, '
'and the graph does not have dgl.NID or dgl.EID columns')
feats[tid, key] = get_storage_func(parent_key, type_).fetch(
column.id_ or default_id, device, pin_memory)
# This class exists to avoid recursion into the feature dictionary returned by the
# prefetcher when calling recursive_apply().
class _PrefetchedGraphFeatures(object):
__slots__ = ['node_feats', 'edge_feats']
def __init__(self, node_feats, edge_feats):
self.node_feats = node_feats
self.edge_feats = edge_feats
def _prefetch_for_subgraph(subg, dataloader):
node_feats, edge_feats = {}, {}
_prefetch_update_feats(
node_feats, subg._node_frames, subg.ntypes, dataloader.graph.get_node_storage,
NID, dataloader.device, dataloader.pin_memory)
_prefetch_update_feats(
edge_feats, subg._edge_frames, subg.canonical_etypes, dataloader.graph.get_edge_storage,
EID, dataloader.device, dataloader.pin_memory)
return _PrefetchedGraphFeatures(node_feats, edge_feats)
def _prefetch_for(item, dataloader):
if isinstance(item, DGLHeteroGraph):
return _prefetch_for_subgraph(item, dataloader)
elif isinstance(item, LazyFeature):
return dataloader.other_storages[item.name].fetch(
item.id_, dataloader.device, dataloader.pin_memory)
else:
exclude_eids = F.cat([eids, F.gather_row(reverse_eid_map, eids)], 0)
return exclude_eids
def _find_exclude_eids_with_reverse_types(g, eids, reverse_etype_map):
exclude_eids = {g.to_canonical_etype(k): v for k, v in eids.items()}
reverse_etype_map = {
g.to_canonical_etype(k): g.to_canonical_etype(v)
for k, v in reverse_etype_map.items()}
exclude_eids.update({reverse_etype_map[k]: v for k, v in exclude_eids.items()})
return exclude_eids
return None
def _find_exclude_eids(g, exclude_mode, eids, **kwargs):
"""Find all edge IDs to exclude according to :attr:`exclude_mode`.
Parameters
----------
g : DGLGraph
The graph.
exclude_mode : str, optional
Can be either of the following,
None (default)
Does not exclude any edge.
'self'
Exclude the given edges themselves but nothing else.
'reverse_id'
Exclude all edges specified in ``eids``, as well as their reverse edges
of the same edge type.
The mapping from each edge ID to its reverse edge ID is specified in
the keyword argument ``reverse_eid_map``.
This mode assumes that the reverse of an edge with ID ``e`` and type
``etype`` will have ID ``reverse_eid_map[e]`` and type ``etype``.
'reverse_types'
Exclude all edges specified in ``eids``, as well as their reverse
edges of the corresponding edge types.
The mapping from each edge type to its reverse edge type is specified
in the keyword argument ``reverse_etype_map``.
This mode assumes that the reverse of an edge with ID ``e`` and type ``etype``
will have ID ``e`` and type ``reverse_etype_map[etype]``.
eids : Tensor or dict[etype, Tensor]
The edge IDs.
reverse_eid_map : Tensor or dict[etype, Tensor]
The mapping from edge ID to its reverse edge ID.
reverse_etype_map : dict[etype, etype]
The mapping from edge etype to its reverse edge type.
"""
if exclude_mode is None:
return None
elif exclude_mode == 'self':
return eids
elif exclude_mode == 'reverse_id':
return _find_exclude_eids_with_reverse_id(g, eids, kwargs['reverse_eid_map'])
elif exclude_mode == 'reverse_types':
return _find_exclude_eids_with_reverse_types(g, eids, kwargs['reverse_etype_map'])
def _await_or_return(x):
if hasattr(x, 'wait'):
return x.wait()
elif isinstance(x, _PrefetchedGraphFeatures):
node_feats = recursive_apply(x.node_feats, _await_or_return)
edge_feats = recursive_apply(x.edge_feats, _await_or_return)
return _PrefetchedGraphFeatures(node_feats, edge_feats)
else:
raise ValueError('unsupported mode {}'.format(exclude_mode))
class Sampler(object):
"""An abstract class that takes in a graph and a set of seed nodes and returns a
structure representing a smaller portion of the graph for computation. It can
be either a list of bipartite graphs (i.e. :class:`BlockSampler`), or a single
subgraph.
return x
def _prefetch(batch, dataloader, stream):
# feats has the same nested structure of batch, except that
# (1) each subgraph is replaced with a pair of node features and edge features, both
# being dictionaries whose keys are (type_id, column_name) and values are either
# tensors or futures.
# (2) each LazyFeature object is replaced with a tensor or future.
# (3) everything else are replaced with None.
#
# Once the futures are fetched, this function waits for them to complete by
# calling its wait() method.
with torch.cuda.stream(stream):
feats = recursive_apply(batch, _prefetch_for, dataloader)
feats = recursive_apply(feats, _await_or_return)
return feats
def _assign_for(item, feat):
if isinstance(item, DGLHeteroGraph):
subg = item
for (tid, key), value in feat.node_feats.items():
assert isinstance(subg._node_frames[tid][key], LazyFeature)
subg._node_frames[tid][key] = value
for (tid, key), value in feat.edge_feats.items():
assert isinstance(subg._edge_frames[tid][key], LazyFeature)
subg._edge_frames[tid][key] = value
return subg
elif isinstance(item, LazyFeature):
return feat
else:
return item
def _prefetcher_entry(dataloader_it, dataloader, queue, num_threads, use_alternate_streams):
# PyTorch will set the number of threads to 1 which slows down pin_memory() calls
# in main process if a prefetching thread is created.
if num_threads is not None:
torch.set_num_threads(num_threads)
if use_alternate_streams:
stream = (
torch.cuda.Stream(device=dataloader.device)
if dataloader.device.type == 'cuda' else None)
else:
stream = None
try:
for batch in dataloader_it:
batch = recursive_apply(batch, restore_parent_storage_columns, dataloader.graph)
feats = _prefetch(batch, dataloader, stream)
queue.put((
# batch will be already in pinned memory as per the behavior of
# PyTorch DataLoader.
recursive_apply(batch, lambda x: x.to(dataloader.device, non_blocking=True)),
feats,
stream.record_event() if stream is not None else None,
None))
queue.put((None, None, None, None))
except: # pylint: disable=bare-except
queue.put((None, None, None, ExceptionWrapper(where='in prefetcher')))
# DGLHeteroGraphs have the semantics of lazy feature slicing with subgraphs. Such behavior depends
# on that DGLHeteroGraph's ndata and edata are maintained by Frames. So to maintain compatibility
# with older code, DGLHeteroGraphs and other graph storages are handled separately: (1)
# DGLHeteroGraphs will preserve the lazy feature slicing for subgraphs. (2) Other graph storages
# will not have lazy feature slicing; all feature slicing will be eager.
def remove_parent_storage_columns(item, g):
"""Removes the storage objects in the given graphs' Frames if it is a sub-frame of the
given parent graph, so that the storages are not serialized during IPC from PyTorch
DataLoader workers.
"""
def __init__(self, output_ctx=None):
self.set_output_context(output_ctx)
def sample(self, g, seed_nodes, exclude_eids=None):
"""Sample a structure from the graph.
Parameters
----------
g : DGLGraph
The original graph.
seed_nodes : Tensor or dict[ntype, Tensor]
The destination nodes by type.
If the graph only has one node type, one can just specify a single tensor
of node IDs.
exclude_eids : Tensor or dict[etype, Tensor]
The edges to exclude from computation dependency.
Returns
-------
Tensor or dict[ntype, Tensor]
The nodes whose input features are required for computing the output
representation of :attr:`seed_nodes`.
any
Any data representing the structure.
"""
raise NotImplementedError
def set_output_context(self, ctx):
"""Set the device the generated block or subgraph will be output to.
This should only be set to a cuda device, when multi-processing is not
used in the dataloader (e.g., num_workers is 0).
Parameters
----------
ctx : DGLContext, default None
The device context the sampled blocks will be stored on. This
should only be a CUDA context if multiprocessing is not used in
the dataloader (e.g., num_workers is 0). If this is None, the
sampled blocks will be stored on the same device as the input
graph.
"""
if ctx is not None:
self.output_device = F.to_backend_ctx(ctx)
else:
self.output_device = None
class BlockSampler(Sampler):
"""Abstract class specifying the neighborhood sampling strategy for DGL data loaders.
The main method for BlockSampler is :meth:`sample`,
which generates a list of message flow graphs (MFGs) for a multi-layer GNN given a set of
seed nodes to have their outputs computed.
The default implementation of :meth:`sample` is
to repeat :attr:`num_layers` times the following procedure from the last layer to the first
layer:
* Obtain a frontier. The frontier is defined as a graph with the same nodes as the
original graph but only the edges involved in message passing on the current layer.
Customizable via :meth:`sample_frontier`.
* Optionally, if the task is link prediction or edge classfication, remove edges
connecting training node pairs. If the graph is undirected, also remove the
reverse edges. This is controlled by the argument :attr:`exclude_eids` in
:meth:`sample` method.
* Convert the frontier into a MFG.
* Optionally assign the IDs of the edges in the original graph selected in the first step
to the MFG, controlled by the argument ``return_eids`` in
:meth:`sample` method.
* Prepend the MFG to the MFG list to be returned.
All subclasses should override :meth:`sample_frontier`
method while specifying the number of layers to sample in :attr:`num_layers` argument.
Parameters
----------
num_layers : int
The number of layers to sample.
return_eids : bool, default False
Whether to return the edge IDs involved in message passing in the MFG.
If True, the edge IDs will be stored as an edge feature named ``dgl.EID``.
output_ctx : DGLContext, default None
The context the sampled blocks will be stored on. This should only be
a CUDA context if multiprocessing is not used in the dataloader (e.g.,
num_workers is 0). If this is None, the sampled blocks will be stored
on the same device as the input graph.
exclude_edges_in_frontier : bool, default False
If True, the :func:`sample_frontier` method will receive an argument
:attr:`exclude_eids` containing the edge IDs from the original graph to exclude.
The :func:`sample_frontier` method must return a graph that does not contain
the edges corresponding to the excluded edges. No additional postprocessing
will be done.
Otherwise, the edges will be removed *after* :func:`sample_frontier` returns.
Notes
-----
For the concept of frontiers and MFGs, please refer to
:ref:`User Guide Section 6 <guide-minibatch>` and
:doc:`Minibatch Training Tutorials <tutorials/large/L0_neighbor_sampling_overview>`.
if not isinstance(item, DGLHeteroGraph) or not isinstance(g, DGLHeteroGraph):
return item
for subframe, frame in zip(
itertools.chain(item._node_frames, item._edge_frames),
itertools.chain(g._node_frames, g._edge_frames)):
for key in list(subframe.keys()):
subcol = subframe._columns[key] # directly get the column object
if isinstance(subcol, LazyFeature):
continue
col = frame._columns.get(key, None)
if col is None:
continue
if col.storage is subcol.storage:
subcol.storage = None
return item
def restore_parent_storage_columns(item, g):
"""Restores the storage objects in the given graphs' Frames if it is a sub-frame of the
given parent graph (i.e. when the storage object is None).
"""
def __init__(self, num_layers, return_eids=False, output_ctx=None):
super().__init__(output_ctx)
self.num_layers = num_layers
self.return_eids = return_eids
# pylint: disable=unused-argument
@staticmethod
def assign_block_eids(block, frontier):
"""Assigns edge IDs from the original graph to the message flow graph (MFG).
See also
--------
BlockSampler
"""
for etype in block.canonical_etypes:
block.edges[etype].data[EID] = frontier.edges[etype].data[EID][
block.edges[etype].data[EID]]
return block
# This is really a hack working around the lack of GPU-based neighbor sampling
# with edge exclusion.
@classmethod
def exclude_edges_in_frontier(cls, g):
"""Returns whether the sampler will exclude edges in :func:`sample_frontier`.
If this method returns True, the method :func:`sample_frontier` will receive an
argument :attr:`exclude_eids` from :func:`sample`. :func:`sample_frontier`
is then responsible for removing those edges.
If this method returns False, :func:`sample` will be responsible for
removing the edges.
When subclassing :class:`BlockSampler`, this method should return True when you
would like to remove the excluded edges in your :func:`sample_frontier` method.
By default this method returns False.
Parameters
----------
g : DGLGraph
The original graph
Returns
-------
bool
Whether :func:`sample_frontier` will receive an argument :attr:`exclude_eids`.
"""
return False
def sample_frontier(self, block_id, g, seed_nodes, exclude_eids=None):
"""Generate the frontier given the destination nodes.
The subclasses should override this function.
Parameters
----------
block_id : int
Represents which GNN layer the frontier is generated for.
g : DGLGraph
The original graph.
seed_nodes : Tensor or dict[ntype, Tensor]
The destination nodes by node type.
If the graph only has one node type, one can just specify a single tensor
of node IDs.
exclude_eids: Tensor or dict
Edge IDs to exclude during sampling neighbors for the seed nodes.
This argument can take a single ID tensor or a dictionary of edge types and ID tensors.
If a single tensor is given, the graph must only have one type of nodes.
Returns
-------
DGLGraph
The frontier generated for the current layer.
Notes
-----
For the concept of frontiers and MFGs, please refer to
:ref:`User Guide Section 6 <guide-minibatch>` and
:doc:`Minibatch Training Tutorials <tutorials/large/L0_neighbor_sampling_overview>`.
"""
raise NotImplementedError
def sample(self, g, seed_nodes, exclude_eids=None):
"""Generate the a list of MFGs given the destination nodes.
Parameters
----------
g : DGLGraph
The original graph.
seed_nodes : Tensor or dict[ntype, Tensor]
The destination nodes by node type.
If the graph only has one node type, one can just specify a single tensor
of node IDs.
exclude_eids : Tensor or dict[etype, Tensor]
The edges to exclude from computation dependency.
Returns
-------
list[DGLGraph]
The MFGs generated for computing the multi-layer GNN output.
Notes
-----
For the concept of frontiers and MFGs, please refer to
:ref:`User Guide Section 6 <guide-minibatch>` and
:doc:`Minibatch Training Tutorials <tutorials/large/L0_neighbor_sampling_overview>`.
"""
blocks = []
if isinstance(g, DistGraph):
# TODO:(nv-dlasalle) dist graphs may not have an associated graph,
# causing an error when trying to fetch the device, so for now,
# always assume the distributed graph's device is CPU.
graph_device = F.cpu()
if not isinstance(item, DGLHeteroGraph) or not isinstance(g, DGLHeteroGraph):
return item
for subframe, frame in zip(
itertools.chain(item._node_frames, item._edge_frames),
itertools.chain(g._node_frames, g._edge_frames)):
for key in subframe.keys():
subcol = subframe._columns[key]
if isinstance(subcol, LazyFeature):
continue
col = frame._columns.get(key, None)
if col is None:
continue
if subcol.storage is None:
subcol.storage = col.storage
return item
class _PrefetchingIter(object):
def __init__(self, dataloader, dataloader_it, use_thread=False, use_alternate_streams=True,
num_threads=None):
self.queue = Queue(1)
self.dataloader_it = dataloader_it
self.dataloader = dataloader
self.graph_sampler = self.dataloader.graph_sampler
self.pin_memory = self.dataloader.pin_memory
self.num_threads = num_threads
self.use_thread = use_thread
self.use_alternate_streams = use_alternate_streams
if use_thread:
thread = threading.Thread(
target=_prefetcher_entry,
args=(dataloader_it, dataloader, self.queue, num_threads, use_alternate_streams),
daemon=True)
thread.start()
self.thread = thread
def __iter__(self):
return self
def _next_non_threaded(self):
batch = next(self.dataloader_it)
batch = recursive_apply(batch, restore_parent_storage_columns, self.dataloader.graph)
device = self.dataloader.device
if self.use_alternate_streams:
stream = torch.cuda.Stream(device=device) if device.type == 'cuda' else None
else:
graph_device = g.device
for block_id in reversed(range(self.num_layers)):
seed_nodes_in = to_device(seed_nodes, graph_device)
if self.exclude_edges_in_frontier(g):
frontier = self.sample_frontier(
block_id, g, seed_nodes_in, exclude_eids=exclude_eids)
else:
frontier = self.sample_frontier(block_id, g, seed_nodes_in)
if self.output_device is not None:
frontier = frontier.to(self.output_device)
seed_nodes_out = to_device(seed_nodes, self.output_device)
else:
seed_nodes_out = seed_nodes
# Removing edges from the frontier for link prediction training falls
# into the category of frontier postprocessing
if not self.exclude_edges_in_frontier(g):
frontier = exclude_edges(frontier, exclude_eids, self.output_device)
block = transform.to_block(frontier, seed_nodes_out)
if self.return_eids:
self.assign_block_eids(block, frontier)
seed_nodes = {ntype: block.srcnodes[ntype].data[NID] for ntype in block.srctypes}
blocks.insert(0, block)
return blocks[0].srcdata[NID], blocks[-1].dstdata[NID], blocks
def sample_blocks(self, g, seed_nodes, exclude_eids=None):
"""Deprecated and identical to :meth:`sample`.
"""
return self.sample(g, seed_nodes, exclude_eids)
class Collator(ABC):
"""Abstract DGL collator for training GNNs on downstream tasks stochastically.
Provides a :attr:`dataset` object containing the collection of all nodes or edges,
as well as a :attr:`collate` method that combines a set of items from
:attr:`dataset` and obtains the message flow graphs (MFGs).
Notes
-----
For the concept of MFGs, please refer to
:ref:`User Guide Section 6 <guide-minibatch>` and
:doc:`Minibatch Training Tutorials <tutorials/large/L0_neighbor_sampling_overview>`.
"""
@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 MFGs.
Parameters
----------
items : list[str, int]
The list of node or edge IDs or type-ID pairs.
Notes
-----
For the concept of MFGs, please refer to
:ref:`User Guide Section 6 <guide-minibatch>` and
:doc:`Minibatch Training Tutorials <tutorials/large/L0_neighbor_sampling_overview>`.
"""
raise NotImplementedError
class NodeCollator(Collator):
"""DGL collator to combine nodes and their computation dependencies within a minibatch for
training node classification or regression on a single graph with neighborhood sampling.
Parameters
----------
g : DGLGraph
The graph.
nids : Tensor or dict[ntype, Tensor]
The node set to compute outputs.
graph_sampler : dgl.dataloading.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.dataloading.MultiLayerNeighborSampler([15, 10, 5])
>>> collator = dgl.dataloading.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)
Notes
-----
For the concept of MFGs, please refer to
:ref:`User Guide Section 6 <guide-minibatch>` and
:doc:`Minibatch Training Tutorials <tutorials/large/L0_neighbor_sampling_overview>`.
stream = None
feats = _prefetch(batch, self.dataloader, stream)
batch = recursive_apply(batch, lambda x: x.to(device, non_blocking=True))
stream_event = stream.record_event() if stream is not None else None
return batch, feats, stream_event
def _next_threaded(self):
batch, feats, stream_event, exception = self.queue.get()
if batch is None:
self.thread.join()
if exception is None:
raise StopIteration
exception.reraise()
return batch, feats, stream_event
def __next__(self):
batch, feats, stream_event = \
self._next_non_threaded() if not self.use_thread else self._next_threaded()
batch = recursive_apply_pair(batch, feats, _assign_for)
if stream_event is not None:
stream_event.wait()
return batch
# Make them classes to work with pickling in mp.spawn
class CollateWrapper(object):
"""Wraps a collate function with :func:`remove_parent_storage_columns` for serializing
from PyTorch DataLoader workers.
"""
def __init__(self, g, nids, graph_sampler):
def __init__(self, sample_func, g):
self.sample_func = sample_func
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.graph_sampler = graph_sampler
self.nids = utils.prepare_tensor_or_dict(g, nids, 'nids')
self._dataset = utils.maybe_flatten_dict(self.nids)
def __call__(self, items):
batch = self.sample_func(self.g, items)
return recursive_apply(batch, remove_parent_storage_columns, self.g)
@property
def dataset(self):
return self._dataset
def collate(self, items):
"""Find the list of MFGs necessary for computing the representation of given
nodes for a node classification/regression task.
Parameters
----------
items : list[int] or list[tuple[str, int]]
Either a list of node IDs (for homogeneous graphs), or a list of node type-ID
pairs (for heterogeneous graphs).
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.
MFGs : list[DGLGraph]
The list of MFGs 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)
items = utils.prepare_tensor_or_dict(self.g, items, 'items')
input_nodes, output_nodes, blocks = self.graph_sampler.sample(self.g, items)
return input_nodes, output_nodes, blocks
class EdgeCollator(Collator):
"""DGL collator to combine edges and their computation dependencies within a minibatch for
training edge classification, edge regression, or link prediction on a single graph
with neighborhood sampling.
Given a set of edges, the collate function will yield
* A tensor of input nodes necessary for computing the representation on edges, or
a dictionary of node type names and such tensors.
* A subgraph that contains only the edges in the minibatch and their incident nodes.
Note that the graph has an identical metagraph with the original graph.
* If a negative sampler is given, another graph that contains the "negative edges",
connecting the source and destination nodes yielded from the given negative sampler.
* A list of MFGs necessary for computing the representation of the incident nodes
of the edges in the minibatch.
Parameters
----------
g : DGLGraph
The graph from which the edges are iterated in minibatches and the subgraphs
are generated.
eids : Tensor or dict[etype, Tensor]
The edge set in graph :attr:`g` to compute outputs.
graph_sampler : dgl.dataloading.BlockSampler
The neighborhood sampler.
g_sampling : DGLGraph, optional
The graph where neighborhood sampling and message passing is performed.
Note that this is not necessarily the same as :attr:`g`.
If None, assume to be the same as :attr:`g`.
exclude : str, optional
Whether and how to exclude dependencies related to the sampled edges in the
minibatch. Possible values are
* None, which excludes nothing.
* ``'self'``, which excludes the sampled edges themselves but nothing else.
* ``'reverse_id'``, which excludes the reverse edges of the sampled edges. The said
reverse edges have the same edge type as the sampled edges. Only works
on edge types whose source node type is the same as its destination node type.
* ``'reverse_types'``, which excludes the reverse edges of the sampled edges. The
said reverse edges have different edge types from the sampled edges.
If ``g_sampling`` is given, ``exclude`` is ignored and will be always ``None``.
reverse_eids : Tensor or dict[etype, Tensor], optional
A tensor of reverse edge ID mapping. The i-th element indicates the ID of
the i-th edge's reverse edge.
If the graph is heterogeneous, this argument requires a dictionary of edge
types and the reverse edge ID mapping tensors.
Required and only used when ``exclude`` is set to ``reverse_id``.
For heterogeneous graph this will be a dict of edge type and edge IDs. Note that
only the edge types whose source node type is the same as destination node type
are needed.
reverse_etypes : dict[etype, etype], optional
The mapping from the edge type to its reverse edge type.
Required and only used when ``exclude`` is set to ``reverse_types``.
negative_sampler : callable, optional
The negative sampler. Can be omitted if no negative sampling is needed.
class WorkerInitWrapper(object):
"""Wraps the :attr:`worker_init_fn` argument of the DataLoader to set the number of DGL
OMP threads to 1 for PyTorch DataLoader workers.
"""
def __init__(self, func):
self.func = func
The negative sampler must be a callable that takes in the following arguments:
def __call__(self, worker_id):
set_num_threads(1)
if self.func is not None:
self.func(worker_id)
* The original (heterogeneous) graph.
* The ID array of sampled edges in the minibatch, or the dictionary of edge
types and ID array of sampled edges in the minibatch if the graph is
heterogeneous.
def create_tensorized_dataset(indices, batch_size, drop_last, use_ddp, ddp_seed):
"""Converts a given indices tensor to a TensorizedDataset, an IterableDataset
that returns views of the original tensor, to reduce overhead from having
a list of scalar tensors in default PyTorch DataLoader implementation.
"""
if use_ddp:
return DDPTensorizedDataset(indices, batch_size, drop_last, ddp_seed)
else:
return TensorizedDataset(indices, batch_size, drop_last)
It should return
* A pair of source and destination node ID arrays as negative samples,
or a dictionary of edge types and such pairs if the graph is heterogenenous.
class DataLoader(torch.utils.data.DataLoader):
"""DataLoader class."""
def __init__(self, graph, indices, graph_sampler, device='cpu', use_ddp=False,
ddp_seed=0, batch_size=1, drop_last=False, shuffle=False,
use_prefetch_thread=False, use_alternate_streams=True, **kwargs):
self.graph = graph
A set of builtin negative samplers are provided in
:ref:`the negative sampling module <api-dataloading-negative-sampling>`.
try:
if isinstance(indices, Mapping):
indices = {k: (torch.tensor(v) if not torch.is_tensor(v) else v)
for k, v in indices.items()}
else:
indices = torch.tensor(indices) if not torch.is_tensor(indices) else indices
except: # pylint: disable=bare-except
# ignore when it fails to convert to torch Tensors.
pass
if (torch.is_tensor(indices) or (
isinstance(indices, Mapping) and
all(torch.is_tensor(v) for v in indices.values()))):
self.dataset = create_tensorized_dataset(
indices, batch_size, drop_last, use_ddp, ddp_seed)
else:
self.dataset = indices
Examples
--------
The following example shows how to train a 3-layer GNN for edge classification on a
set of edges ``train_eid`` on a homogeneous undirected graph. Each node takes
messages from all neighbors.
Say that you have an array of source node IDs ``src`` and another array of destination
node IDs ``dst``. One can make it bidirectional by adding another set of edges
that connects from ``dst`` to ``src``:
>>> g = dgl.graph((torch.cat([src, dst]), torch.cat([dst, src])))
One can then know that the ID difference of an edge and its reverse edge is ``|E|``,
where ``|E|`` is the length of your source/destination array. The reverse edge
mapping can be obtained by
>>> E = len(src)
>>> reverse_eids = torch.cat([torch.arange(E, 2 * E), torch.arange(0, E)])
Note that the sampled edges as well as their reverse edges are removed from
computation dependencies of the incident nodes. This is a common trick to avoid
information leakage.
>>> sampler = dgl.dataloading.MultiLayerNeighborSampler([15, 10, 5])
>>> collator = dgl.dataloading.EdgeCollator(
... g, train_eid, sampler, exclude='reverse_id',
... reverse_eids=reverse_eids)
>>> 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, pair_graph, blocks in dataloader:
... train_on(input_nodes, pair_graph, blocks)
To train a 3-layer GNN for link prediction on a set of edges ``train_eid`` on a
homogeneous graph where each node takes messages from all neighbors (assume the
backend is PyTorch), with 5 uniformly chosen negative samples per edge:
>>> sampler = dgl.dataloading.MultiLayerNeighborSampler([15, 10, 5])
>>> neg_sampler = dgl.dataloading.negative_sampler.Uniform(5)
>>> collator = dgl.dataloading.EdgeCollator(
... g, train_eid, sampler, exclude='reverse_id',
... reverse_eids=reverse_eids, negative_sampler=neg_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, pos_pair_graph, neg_pair_graph, blocks in dataloader:
... train_on(input_nodse, pair_graph, neg_pair_graph, blocks)
For heterogeneous graphs, the reverse of an edge may have a different edge type
from the original edge. For instance, consider that you have an array of
user-item clicks, representated by a user array ``user`` and an item array ``item``.
You may want to build a heterogeneous graph with a user-click-item relation and an
item-clicked-by-user relation.
>>> g = dgl.heterograph({
... ('user', 'click', 'item'): (user, item),
... ('item', 'clicked-by', 'user'): (item, user)})
To train a 3-layer GNN for edge classification on a set of edges ``train_eid`` with
type ``click``, you can write
>>> sampler = dgl.dataloading.MultiLayerNeighborSampler([15, 10, 5])
>>> collator = dgl.dataloading.EdgeCollator(
... g, {'click': train_eid}, sampler, exclude='reverse_types',
... reverse_etypes={'click': 'clicked-by', 'clicked-by': 'click'})
>>> 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, pair_graph, blocks in dataloader:
... train_on(input_nodes, pair_graph, blocks)
To train a 3-layer GNN for link prediction on a set of edges ``train_eid`` with type
``click``, you can write
>>> sampler = dgl.dataloading.MultiLayerNeighborSampler([15, 10, 5])
>>> neg_sampler = dgl.dataloading.negative_sampler.Uniform(5)
>>> collator = dgl.dataloading.EdgeCollator(
... g, train_eid, sampler, exclude='reverse_types',
... reverse_etypes={'click': 'clicked-by', 'clicked-by': 'click'},
... negative_sampler=neg_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, pos_pair_graph, neg_pair_graph, blocks in dataloader:
... train_on(input_nodes, pair_graph, neg_pair_graph, blocks)
Notes
-----
For the concept of MFGs, please refer to
:ref:`User Guide Section 6 <guide-minibatch>` and
:doc:`Minibatch Training Tutorials <tutorials/large/L0_neighbor_sampling_overview>`.
"""
def __init__(self, g, eids, graph_sampler, g_sampling=None, exclude=None,
reverse_eids=None, reverse_etypes=None, negative_sampler=None):
self.g = g
if not isinstance(eids, Mapping):
assert len(g.etypes) == 1, \
"eids should be a dict of etype and ids for graph with multiple etypes"
self.ddp_seed = ddp_seed
self._shuffle_dataset = shuffle
self.graph_sampler = graph_sampler
# One may wish to iterate over the edges in one graph while perform sampling in
# another graph. This may be the case for iterating over validation and test
# edge set while perform neighborhood sampling on the graph formed by only
# the training edge set.
# See GCMC for an example usage.
self.device = torch.device(device)
self.use_alternate_streams = use_alternate_streams
if self.device.type == 'cuda' and self.device.index is None:
self.device = torch.device('cuda', torch.cuda.current_device())
self.use_prefetch_thread = use_prefetch_thread
worker_init_fn = WorkerInitWrapper(kwargs.get('worker_init_fn', None))
# Instantiate all the formats if the number of workers is greater than 0.
if kwargs.get('num_workers', 0) > 0 and hasattr(self.graph, 'create_formats_'):
self.graph.create_formats_()
self.other_storages = {}
super().__init__(
self.dataset,
collate_fn=CollateWrapper(self.graph_sampler.sample, graph),
batch_size=None,
worker_init_fn=worker_init_fn,
**kwargs)
def __iter__(self):
if self._shuffle_dataset:
self.dataset.shuffle()
# When using multiprocessing PyTorch sometimes set the number of PyTorch threads to 1
# when spawning new Python threads. This drastically slows down pinning features.
num_threads = torch.get_num_threads() if self.num_workers > 0 else None
return _PrefetchingIter(
self, super().__iter__(), use_thread=self.use_prefetch_thread,
use_alternate_streams=self.use_alternate_streams, num_threads=num_threads)
# To allow data other than node/edge data to be prefetched.
def attach_data(self, name, data):
"""Add a data other than node and edge features for prefetching."""
self.other_storages[name] = wrap_storage(data)
# Alias
class NodeDataLoader(DataLoader):
"""NodeDataLoader class."""
class EdgeDataLoader(DataLoader):
"""EdgeDataLoader class."""
def __init__(self, graph, indices, graph_sampler, device='cpu', use_ddp=False,
ddp_seed=0, batch_size=1, drop_last=False, shuffle=False,
use_prefetch_thread=False, use_alternate_streams=True,
exclude=None, reverse_eids=None, reverse_etypes=None, negative_sampler=None,
g_sampling=None, **kwargs):
if g_sampling is not None:
self.g_sampling = g_sampling
self.exclude = None
else:
self.g_sampling = self.g
self.exclude = exclude
self.reverse_eids = reverse_eids
self.reverse_etypes = reverse_etypes
self.negative_sampler = negative_sampler
self.eids = utils.prepare_tensor_or_dict(g, eids, 'eids')
self._dataset = utils.maybe_flatten_dict(self.eids)
@property
def dataset(self):
return self._dataset
def _collate(self, items):
if isinstance(items[0], tuple):
# returns a list of pairs: group them by node types into a dict
items = utils.group_as_dict(items)
items = utils.prepare_tensor_or_dict(self.g_sampling, items, 'items')
pair_graph = self.g.edge_subgraph(items)
seed_nodes = pair_graph.ndata[NID]
exclude_eids = _find_exclude_eids(
self.g_sampling,
self.exclude,
items,
reverse_eid_map=self.reverse_eids,
reverse_etype_map=self.reverse_etypes)
input_nodes, _, blocks = self.graph_sampler.sample(
self.g_sampling, seed_nodes, exclude_eids=exclude_eids)
return input_nodes, pair_graph, blocks
def _collate_with_negative_sampling(self, items):
if isinstance(items[0], tuple):
# returns a list of pairs: group them by node types into a dict
items = utils.group_as_dict(items)
items = utils.prepare_tensor_or_dict(self.g_sampling, items, 'items')
pair_graph = self.g.edge_subgraph(items, relabel_nodes=False)
induced_edges = pair_graph.edata[EID]
neg_srcdst = self.negative_sampler(self.g, items)
if not isinstance(neg_srcdst, Mapping):
assert len(self.g.etypes) == 1, \
'graph has multiple or no edge types; '\
'please return a dict in negative sampler.'
neg_srcdst = {self.g.canonical_etypes[0]: neg_srcdst}
# Get dtype from a tuple of tensors
dtype = F.dtype(list(neg_srcdst.values())[0][0])
ctx = F.context(pair_graph)
neg_edges = {
etype: neg_srcdst.get(etype, (F.copy_to(F.tensor([], dtype), ctx),
F.copy_to(F.tensor([], dtype), ctx)))
for etype in self.g.canonical_etypes}
neg_pair_graph = heterograph(
neg_edges, {ntype: self.g.number_of_nodes(ntype) for ntype in self.g.ntypes})
pair_graph, neg_pair_graph = transform.compact_graphs([pair_graph, neg_pair_graph])
pair_graph.edata[EID] = induced_edges
seed_nodes = pair_graph.ndata[NID]
exclude_eids = _find_exclude_eids(
self.g_sampling,
self.exclude,
items,
reverse_eid_map=self.reverse_eids,
reverse_etype_map=self.reverse_etypes)
input_nodes, _, blocks = self.graph_sampler.sample(
self.g_sampling, seed_nodes, exclude_eids=exclude_eids)
return input_nodes, pair_graph, neg_pair_graph, blocks
def collate(self, items):
"""Combines the sampled edges into a minibatch for edge classification, edge
regression, and link prediction tasks.
Parameters
----------
items : list[int] or list[tuple[str, int]]
Either a list of edge IDs (for homogeneous graphs), or a list of edge type-ID
pairs (for heterogeneous graphs).
Returns
-------
Either ``(input_nodes, pair_graph, blocks)``, or
``(input_nodes, pair_graph, negative_pair_graph, blocks)`` if negative sampling is
enabled.
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.
pair_graph : DGLGraph
The graph that contains only the edges in the minibatch as well as their incident
nodes.
Note that the metagraph of this graph will be identical to that of the original
graph.
negative_pair_graph : DGLGraph
The graph that contains only the edges connecting the source and destination nodes
yielded from the given negative sampler, if negative sampling is enabled.
Note that the metagraph of this graph will be identical to that of the original
graph.
blocks : list[DGLGraph]
The list of MFGs necessary for computing the representation of the edges.
"""
if self.negative_sampler is None:
return self._collate(items)
else:
return self._collate_with_negative_sampling(items)
dgl_warning(
"g_sampling is deprecated. "
"Please merge g_sampling and the original graph into one graph and use "
"the exclude argument to specify which edges you don't want to sample.")
if isinstance(graph_sampler, BlockSampler):
graph_sampler = EdgeBlockSampler(
graph_sampler, exclude=exclude, reverse_eids=reverse_eids,
reverse_etypes=reverse_etypes, negative_sampler=negative_sampler)
super().__init__(
graph, indices, graph_sampler, device=device, use_ddp=use_ddp, ddp_seed=ddp_seed,
batch_size=batch_size, drop_last=drop_last, shuffle=shuffle,
use_prefetch_thread=use_prefetch_thread, use_alternate_streams=use_alternate_streams,
**kwargs)
######## Graph DataLoaders ########
# GraphDataLoader loads a set of graphs so it's not relevant to the above. They are currently
# copied from the old DataLoader implementation.
PYTORCH_VER = LooseVersion(torch.__version__)
PYTORCH_16 = PYTORCH_VER >= LooseVersion("1.6.0")
PYTORCH_17 = PYTORCH_VER >= LooseVersion("1.7.0")
def _create_dist_sampler(dataset, dataloader_kwargs, ddp_seed):
# Note: will change the content of dataloader_kwargs
dist_sampler_kwargs = {'shuffle': dataloader_kwargs['shuffle']}
dataloader_kwargs['shuffle'] = False
if PYTORCH_16:
dist_sampler_kwargs['seed'] = ddp_seed
if PYTORCH_17:
dist_sampler_kwargs['drop_last'] = dataloader_kwargs['drop_last']
dataloader_kwargs['drop_last'] = False
return DistributedSampler(dataset, **dist_sampler_kwargs)
class GraphCollator(object):
"""Given a set of graphs as well as their graph-level data, the collate function will batch the
......@@ -939,8 +641,8 @@ class GraphCollator(object):
"""
elem = items[0]
elem_type = type(elem)
if isinstance(elem, DGLGraph):
batched_graphs = batch(items)
if isinstance(elem, DGLHeteroGraph):
batched_graphs = batch_graphs(items)
return batched_graphs
elif F.is_tensor(elem):
return F.stack(items, 0)
......@@ -975,8 +677,89 @@ class GraphCollator(object):
raise TypeError(self.graph_collate_err_msg_format.format(elem_type))
class SubgraphIterator(object):
"""Abstract class representing an iterator that yields a subgraph given a graph.
class GraphDataLoader(torch.utils.data.DataLoader):
"""PyTorch dataloader for batch-iterating over a set of graphs, generating the batched
graph and corresponding label tensor (if provided) of the said minibatch.
Parameters
----------
collate_fn : Function, default is None
The customized collate function. Will use the default collate
function if not given.
use_ddp : boolean, optional
If True, tells the DataLoader to split the training set for each
participating process appropriately using
:class:`torch.utils.data.distributed.DistributedSampler`.
Overrides the :attr:`sampler` argument of :class:`torch.utils.data.DataLoader`.
ddp_seed : int, optional
The seed for shuffling the dataset in
:class:`torch.utils.data.distributed.DistributedSampler`.
Only effective when :attr:`use_ddp` is True.
kwargs : dict
Arguments being passed to :py:class:`torch.utils.data.DataLoader`.
Examples
--------
To train a GNN for graph classification on a set of graphs in ``dataset`` (assume
the backend is PyTorch):
>>> dataloader = dgl.dataloading.GraphDataLoader(
... dataset, batch_size=1024, shuffle=True, drop_last=False, num_workers=4)
>>> for batched_graph, labels in dataloader:
... train_on(batched_graph, labels)
**Using with Distributed Data Parallel**
If you are using PyTorch's distributed training (e.g. when using
:mod:`torch.nn.parallel.DistributedDataParallel`), you can train the model by
turning on the :attr:`use_ddp` option:
>>> dataloader = dgl.dataloading.GraphDataLoader(
... dataset, use_ddp=True, batch_size=1024, shuffle=True, drop_last=False, num_workers=4)
>>> for epoch in range(start_epoch, n_epochs):
... dataloader.set_epoch(epoch)
... for batched_graph, labels in dataloader:
... train_on(batched_graph, labels)
"""
def __init__(self, g):
self.g = g
collator_arglist = inspect.getfullargspec(GraphCollator).args
def __init__(self, dataset, collate_fn=None, use_ddp=False, ddp_seed=0, **kwargs):
collator_kwargs = {}
dataloader_kwargs = {}
for k, v in kwargs.items():
if k in self.collator_arglist:
collator_kwargs[k] = v
else:
dataloader_kwargs[k] = v
if collate_fn is None:
self.collate = GraphCollator(**collator_kwargs).collate
else:
self.collate = collate_fn
self.use_ddp = use_ddp
if use_ddp:
self.dist_sampler = _create_dist_sampler(dataset, dataloader_kwargs, ddp_seed)
dataloader_kwargs['sampler'] = self.dist_sampler
super().__init__(dataset=dataset, collate_fn=self.collate, **dataloader_kwargs)
def set_epoch(self, epoch):
"""Sets the epoch number for the underlying sampler which ensures all replicas
to use a different ordering for each epoch.
Only available when :attr:`use_ddp` is True.
Calls :meth:`torch.utils.data.distributed.DistributedSampler.set_epoch`.
Parameters
----------
epoch : int
The epoch number.
"""
if self.use_ddp:
self.dist_sampler.set_epoch(epoch)
else:
raise DGLError('set_epoch is only available when use_ddp is True.')
python/dgl/dataloading/dist_dataloader.py 0 → 100644
View file @ 701b4fcc
"""Distributed dataloaders.
"""
import inspect
from ..distributed import DistDataLoader
# Still depends on the legacy NodeCollator...
from .._dataloading.dataloader import NodeCollator, EdgeCollator
def _remove_kwargs_dist(kwargs):
if 'num_workers' in kwargs:
del kwargs['num_workers']
if 'pin_memory' in kwargs:
del kwargs['pin_memory']
print('Distributed DataLoaders do not support pin_memory.')
return kwargs
class DistNodeDataLoader(DistDataLoader):
"""PyTorch dataloader for batch-iterating over a set of nodes, generating the list
of message flow graphs (MFGs) as computation dependency of the said minibatch, on
a distributed graph.
All the arguments have the same meaning as the single-machine counterpart
:class:`dgl.dataloading.pytorch.NodeDataLoader` except the first argument
:attr:`g` which must be a :class:`dgl.distributed.DistGraph`.
Parameters
----------
g : DistGraph
The distributed graph.
nids, graph_sampler, device, kwargs :
See :class:`dgl.dataloading.pytorch.NodeDataLoader`.
See also
--------
dgl.dataloading.pytorch.NodeDataLoader
"""
def __init__(self, g, nids, graph_sampler, device=None, **kwargs):
collator_kwargs = {}
dataloader_kwargs = {}
_collator_arglist = inspect.getfullargspec(NodeCollator).args
for k, v in kwargs.items():
if k in _collator_arglist:
collator_kwargs[k] = v
else:
dataloader_kwargs[k] = v
if device is None:
# for the distributed case default to the CPU
device = 'cpu'
assert device == 'cpu', 'Only cpu is supported in the case of a DistGraph.'
# Distributed DataLoader currently does not support heterogeneous graphs
# and does not copy features. Fallback to normal solution
self.collator = NodeCollator(g, nids, graph_sampler, **collator_kwargs)
_remove_kwargs_dist(dataloader_kwargs)
super().__init__(self.collator.dataset,
collate_fn=self.collator.collate,
**dataloader_kwargs)
self.device = device
class DistEdgeDataLoader(DistDataLoader):
"""PyTorch dataloader for batch-iterating over a set of edges, generating the list
of message flow graphs (MFGs) as computation dependency of the said minibatch for
edge classification, edge regression, and link prediction, on a distributed
graph.
All the arguments have the same meaning as the single-machine counterpart
:class:`dgl.dataloading.pytorch.EdgeDataLoader` except the first argument
:attr:`g` which must be a :class:`dgl.distributed.DistGraph`.
Parameters
----------
g : DistGraph
The distributed graph.
eids, graph_sampler, device, kwargs :
See :class:`dgl.dataloading.pytorch.EdgeDataLoader`.
See also
--------
dgl.dataloading.pytorch.EdgeDataLoader
"""
def __init__(self, g, eids, graph_sampler, device=None, **kwargs):
collator_kwargs = {}
dataloader_kwargs = {}
_collator_arglist = inspect.getfullargspec(EdgeCollator).args
for k, v in kwargs.items():
if k in _collator_arglist:
collator_kwargs[k] = v
else:
dataloader_kwargs[k] = v
if device is None:
# for the distributed case default to the CPU
device = 'cpu'
assert device == 'cpu', 'Only cpu is supported in the case of a DistGraph.'
# Distributed DataLoader currently does not support heterogeneous graphs
# and does not copy features. Fallback to normal solution
self.collator = EdgeCollator(g, eids, graph_sampler, **collator_kwargs)
_remove_kwargs_dist(dataloader_kwargs)
super().__init__(self.collator.dataset,
collate_fn=self.collator.collate,
**dataloader_kwargs)
self.device = device
python/dgl/dataloading/negative_sampler.py
View file @ 701b4fcc
"""Negative samplers"""
from collections.abc import Mapping
from .. import backend as F
from ..sampling import global_uniform_negative_sampling
class _BaseNegativeSampler(object):
def _generate(self, g, eids, canonical_etype):
......@@ -26,7 +25,7 @@ class _BaseNegativeSampler(object):
eids = {g.to_canonical_etype(k): v for k, v in eids.items()}
neg_pair = {k: self._generate(g, v, k) for k, v in eids.items()}
else:
assert len(g.etypes) == 1, \
assert len(g.canonical_etypes) == 1, \
'please specify a dict of etypes and ids for graphs with multiple edge types'
neg_pair = self._generate(g, eids, g.canonical_etypes[0])
......@@ -64,7 +63,7 @@ class PerSourceUniform(_BaseNegativeSampler):
shape = (shape[0] * self.k,)
src, _ = g.find_edges(eids, etype=canonical_etype)
src = F.repeat(src, self.k, 0)
dst = F.randint(shape, dtype, ctx, 0, g.number_of_nodes(vtype))
dst = F.randint(shape, dtype, ctx, 0, g.num_nodes(vtype))
return src, dst
# Alias
......@@ -90,14 +89,6 @@ class GlobalUniform(_BaseNegativeSampler):
replace : bool, optional
Whether to sample with replacement. Setting it to True will make things
faster. (Default: True)
redundancy : float, optional
Indicates how much more negative samples to actually generate during rejection sampling
before finding the unique pairs.
Increasing it will increase the likelihood of getting :attr:`k` negative samples
per edge, but will also take more time and memory.
(Default: automatically determined by the density of graph)
Notes
-----
......@@ -113,13 +104,11 @@ class GlobalUniform(_BaseNegativeSampler):
>>> neg_sampler(g, torch.LongTensor([0, 1]))
(tensor([0, 1, 3, 2]), tensor([2, 0, 2, 1]))
"""
def __init__(self, k, exclude_self_loops=True, replace=False, redundancy=None):
def __init__(self, k, exclude_self_loops=True, replace=False):
self.k = k
self.exclude_self_loops = exclude_self_loops
self.replace = replace
self.redundancy = redundancy
def _generate(self, g, eids, canonical_etype):
return global_uniform_negative_sampling(
g, len(eids) * self.k, self.exclude_self_loops, self.replace,
canonical_etype, self.redundancy)
return g.global_uniform_negative_sampling(
len(eids) * self.k, self.exclude_self_loops, self.replace, canonical_etype)
python/dgl/dataloading/neighbor_sampler.py 0 → 100644
View file @ 701b4fcc
"""Data loading components for neighbor sampling"""
from ..base import NID, EID
from ..transform import to_block
from .base import BlockSampler
class NeighborSampler(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]]
List of neighbors to sample per edge type for each GNN layer, with the i-th
element being the fanout for the i-th GNN layer.
If only a single integer is provided, DGL assumes that every edge type
will have the same fanout.
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 False
Whether to sample with replacement
prob : str, optional
If given, the probability of each neighbor being sampled is proportional
to the edge feature value with the given name in ``g.edata``. The feature must be
a scalar on each edge.
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 5, 10, 15 neighbors for
the first, second, and third layer respectively (assuming the backend is PyTorch):
>>> sampler = dgl.dataloading.NeighborSampler([5, 10, 15])
>>> dataloader = dgl.dataloading.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(blocks)
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.dataloading.NeighborSampler([
... {('user', 'follows', 'user'): 5,
... ('user', 'plays', 'game'): 4,
... ('game', 'played-by', 'user'): 3}] * 3)
If you would like non-uniform neighbor sampling:
>>> g.edata['p'] = torch.rand(g.num_edges()) # any non-negative 1D vector works
>>> sampler = dgl.dataloading.NeighborSampler([5, 10, 15], prob='p')
Notes
-----
For the concept of MFGs, please refer to
:ref:`User Guide Section 6 <guide-minibatch>` and
:doc:`Minibatch Training Tutorials <tutorials/large/L0_neighbor_sampling_overview>`.
"""
def __init__(self, fanouts, edge_dir='in', prob=None, replace=False, **kwargs):
super().__init__(**kwargs)
self.fanouts = fanouts
self.edge_dir = edge_dir
self.prob = prob
self.replace = replace
def sample_blocks(self, g, seed_nodes, exclude_eids=None):
output_nodes = seed_nodes
blocks = []
for fanout in reversed(self.fanouts):
frontier = g.sample_neighbors(
seed_nodes, fanout, edge_dir=self.edge_dir, prob=self.prob,
replace=self.replace, output_device=self.output_device,
exclude_edges=exclude_eids)
eid = frontier.edata[EID]
block = to_block(frontier, seed_nodes)
block.edata[EID] = eid
seed_nodes = block.srcdata[NID]
blocks.insert(0, block)
return seed_nodes, output_nodes, blocks
MultiLayerNeighborSampler = NeighborSampler
class MultiLayerFullNeighborSampler(NeighborSampler):
"""Sampler that builds computational dependency of node representations by taking messages
from all neighbors for multilayer GNN.
This sampler will make every node gather messages from every single neighbor per edge type.
Parameters
----------
n_layers : int
The number of GNN layers to sample.
return_eids : bool, default False
Whether to return the edge IDs involved in message passing in the MFG.
If True, the edge IDs will be stored as an edge feature named ``dgl.EID``.
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 for the first,
second, and third layer respectively (assuming the backend is PyTorch):
>>> sampler = dgl.dataloading.MultiLayerFullNeighborSampler(3)
>>> dataloader = dgl.dataloading.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(blocks)
Notes
-----
For the concept of MFGs, please refer to
:ref:`User Guide Section 6 <guide-minibatch>` and
:doc:`Minibatch Training Tutorials <tutorials/large/L0_neighbor_sampling_overview>`.
"""
def __init__(self, num_layers, edge_dir='in', prob=None, replace=False, **kwargs):
super().__init__([-1] * num_layers, edge_dir=edge_dir, prob=prob, replace=replace,
**kwargs)
python/dgl/dataloading/shadow.py
View file @ 701b4fcc
"""ShaDow-GNN subgraph samplers."""
from ..utils import prepare_tensor_or_dict
from ..base import NID
from ..sampling.utils import EidExcluder
from .. import transform
from ..sampling import sample_neighbors
from .neighbor import NeighborSamplingMixin
from .dataloader import exclude_edges, Sampler
from ..base import NID
from .base import set_node_lazy_features, set_edge_lazy_features
class ShaDowKHopSampler(NeighborSamplingMixin, Sampler):
class ShaDowKHopSampler(object):
"""K-hop subgraph sampler used by
`ShaDow-GNN <https://arxiv.org/abs/2012.01380>`__.
......@@ -70,29 +68,32 @@ class ShaDowKHopSampler(NeighborSamplingMixin, Sampler):
If you would like non-uniform neighbor sampling:
>>> g.edata['p'] = torch.rand(g.num_edges()) # any non-negative 1D vector works
>>> sampler = dgl.dataloading.MultiLayerNeighborSampler([5, 10, 15], prob='p')
>>> sampler = dgl.dataloading.ShaDowKHopSampler([5, 10, 15], prob='p')
"""
def __init__(self, fanouts, replace=False, prob=None, output_ctx=None):
super().__init__(output_ctx)
def __init__(self, fanouts, replace=False, prob=None, prefetch_node_feats=None,
prefetch_edge_feats=None, output_device=None):
self.fanouts = fanouts
self.replace = replace
self.prob = prob
self.set_output_context(output_ctx)
self.prefetch_node_feats = prefetch_node_feats
self.prefetch_edge_feats = prefetch_edge_feats
self.output_device = output_device
def sample(self, g, seed_nodes, exclude_eids=None):
self._build_fanout(len(self.fanouts), g)
self._build_prob_arrays(g)
seed_nodes = prepare_tensor_or_dict(g, seed_nodes, 'seed nodes')
def sample(self, g, seed_nodes, exclude_edges=None):
"""Sample a subgraph given a tensor of seed nodes."""
output_nodes = seed_nodes
for i in range(len(self.fanouts)):
fanout = self.fanouts[i]
frontier = sample_neighbors(
g, seed_nodes, fanout, replace=self.replace, prob=self.prob_arrays)
for fanout in reversed(self.fanouts):
frontier = g.sample_neighbors(
seed_nodes, fanout, output_device=self.output_device,
replace=self.replace, prob=self.prob, exclude_edges=exclude_edges)
block = transform.to_block(frontier, seed_nodes)
seed_nodes = block.srcdata[NID]
subg = g.subgraph(seed_nodes, relabel_nodes=True)
subg = exclude_edges(subg, exclude_eids, self.output_device)
subg = g.subgraph(seed_nodes, relabel_nodes=True, output_device=self.output_device)
if exclude_edges is not None:
subg = EidExcluder(exclude_edges)(subg)
set_node_lazy_features(subg, self.prefetch_node_feats)
set_edge_lazy_features(subg, self.prefetch_edge_feats)
return seed_nodes, output_nodes, [subg]
return seed_nodes, output_nodes, subg
python/dgl/distributed/dist_graph.py
View file @ 701b4fcc
......@@ -26,6 +26,7 @@ from . import rpc
from . import role
from .server_state import ServerState
from .rpc_server import start_server
from . import graph_services
from .graph_services import find_edges as dist_find_edges
from .graph_services import out_degrees as dist_out_degrees
from .graph_services import in_degrees as dist_in_degrees
......@@ -1223,6 +1224,20 @@ class DistGraph:
'''
self._client.barrier()
def sample_neighbors(self, seed_nodes, fanout, edge_dir='in', prob=None,
exclude_edges=None, replace=False,
output_device=None):
# pylint: disable=unused-argument
"""Sample neighbors from a distributed graph."""
# Currently prob, exclude_edges, output_device, and edge_dir are ignored.
if len(self.etypes) > 1:
frontier = graph_services.sample_etype_neighbors(
self, seed_nodes, ETYPE, fanout, replace=replace)
else:
frontier = graph_services.sample_neighbors(
self, seed_nodes, fanout, replace=replace)
return frontier
def _get_ndata_names(self, ntype=None):
''' Get the names of all node data.
'''
......
python/dgl/frame.py
View file @ 701b4fcc
......@@ -7,6 +7,7 @@ from collections.abc import MutableMapping
from . import backend as F
from .base import DGLError, dgl_warning
from .init import zero_initializer
from .storages import TensorStorage
class _LazyIndex(object):
def __init__(self, index):
......@@ -38,6 +39,23 @@ class _LazyIndex(object):
flat_index = F.gather_row(flat_index, index)
return flat_index
class LazyFeature(object):
"""Placeholder for prefetching from DataLoader.
"""
__slots__ = ['name', 'id_']
def __init__(self, name=None, id_=None):
self.name = name
self.id_ = id_
def to(self, *args, **kwargs): # pylint: disable=invalid-name, unused-argument
"""No-op. For compatibility of :meth:`Frame.to` method."""
return self
@property
def data(self):
"""No-op. For compatibility of :meth:`Frame.__repr__` method."""
return self
class Scheme(namedtuple('Scheme', ['shape', 'dtype'])):
"""The column scheme.
......@@ -77,7 +95,7 @@ def infer_scheme(tensor):
"""
return Scheme(tuple(F.shape(tensor)[1:]), F.dtype(tensor))
class Column(object):
class Column(TensorStorage):
"""A column is a compact store of features of multiple nodes/edges.
It batches all the feature tensors together along the first dimension
......@@ -120,7 +138,7 @@ class Column(object):
Index tensor
"""
def __init__(self, storage, scheme=None, index=None, device=None):
self.storage = storage
super().__init__(storage)
self.scheme = scheme if scheme else infer_scheme(storage)
self.index = index
self.device = device
......@@ -336,10 +354,13 @@ class Frame(MutableMapping):
assert not isinstance(data, Frame) # sanity check for code refactor
# Note that we always create a new column for the given data.
# This avoids two frames accidentally sharing the same column.
self._columns = {k : Column.create(v) for k, v in data.items()}
self._columns = {k : v if isinstance(v, LazyFeature) else Column.create(v)
for k, v in data.items()}
self._num_rows = num_rows
# infer num_rows & sanity check
for name, col in self._columns.items():
if isinstance(col, LazyFeature):
continue
if self._num_rows is None:
self._num_rows = len(col)
elif len(col) != self._num_rows:
......@@ -504,6 +525,10 @@ class Frame(MutableMapping):
data : Column or data convertible to Column
The column data.
"""
if isinstance(data, LazyFeature):
self._columns[name] = data
return
col = Column.create(data)
if len(col) != self.num_rows:
raise DGLError('Expected data to have %d rows, got %d.' %
......
python/dgl/geometry/capi.py
View file @ 701b4fcc
"""Python interfaces to DGL farthest point sampler."""
from dgl._ffi.base import DGLError
import numpy as np
from .._ffi.base import DGLError
from .._ffi.function import _init_api
from .. import backend as F
from .. import ndarray as nd
......
python/dgl/heterograph.py
View file @ 701b4fcc
......@@ -1600,6 +1600,14 @@ class DGLHeteroGraph(object):
# View
#################################################################
def get_node_storage(self, key, ntype=None):
"""Get storage object of node feature of type :attr:`ntype` and name :attr:`key`."""
return self._node_frames[self.get_ntype_id(ntype)]._columns[key]
def get_edge_storage(self, key, etype=None):
"""Get storage object of edge feature of type :attr:`etype` and name :attr:`key`."""
return self._edge_frames[self.get_etype_id(etype)]._columns[key]
@property
def nodes(self):
"""Return a node view
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment