wehub/dmlc--dgl
1
0
Fork 0
Code Issues Pull Requests Actions 1 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:35:51 +08:00
commit c36a561cd8
2172 changed files with 455595 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/dgl/dataloading/__init__.py
+14
View File
@@ -0,0 +1,14 @@
"""Package for dataloaders and samplers."""
from .. import backend as F
from . import negative_sampler
from .base import *
from .cluster_gcn import *
from .graphsaint import *
from .labor_sampler import *
from .neighbor_sampler import *
from .shadow import *
if F.get_preferred_backend() == "pytorch":
from .spot_target import *
from .dataloader import *
python/dgl/dataloading/base.py
+658
View File
@@ -0,0 +1,658 @@
"""Base classes and functionalities for dataloaders"""
import inspect
from collections.abc import Mapping
from .. import backend as F
from ..base import EID, NID
from ..convert import heterograph
from ..frame import LazyFeature
from ..transforms import compact_graphs
from ..utils import context_of, 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):
"""Assign lazy features to the ``ndata`` of the input graph for prefetching optimization.
When used in a :class:`~dgl.dataloading.Sampler`, lazy features mark which data
should be fetched before computation in model. See :ref:`guide-minibatch-prefetching`
for a detailed explanation.
If the graph is homogeneous, this is equivalent to:
.. code:: python
g.ndata.update({k: LazyFeature(k, g.ndata[dgl.NID]) for k in feature_names})
If the graph is heterogeneous, this is equivalent to:
.. code:: python
for type_, names in feature_names.items():
g.nodes[type_].data.update(
{k: LazyFeature(k, g.nodes[type_].data[dgl.NID]) for k in names})
Parameters
----------
g : DGLGraph
The graph.
feature_names : list[str] or dict[str, list[str]]
The feature names to prefetch.
See also
--------
dgl.LazyFeature
"""
return _set_lazy_features(g.nodes, g.ndata, feature_names)
def set_edge_lazy_features(g, feature_names):
"""Assign lazy features to the ``edata`` of the input graph for prefetching optimization.
When used in a :class:`~dgl.dataloading.Sampler`, lazy features mark which data
should be fetched before computation in model. See :ref:`guide-minibatch-prefetching`
for a detailed explanation.
If the graph is homogeneous, this is equivalent to:
.. code:: python
g.edata.update({k: LazyFeature(k, g.edata[dgl.EID]) for k in feature_names})
If the graph is heterogeneous, this is equivalent to:
.. code:: python
for type_, names in feature_names.items():
g.edges[type_].data.update(
{k: LazyFeature(k, g.edges[type_].data[dgl.EID]) for k in names})
Parameters
----------
g : DGLGraph
The graph.
feature_names : list[str] or dict[etype, list[str]]
The feature names to prefetch. The ``etype`` key is either a string
or a triplet.
See also
--------
dgl.LazyFeature
"""
return _set_lazy_features(g.edges, g.edata, feature_names)
def set_src_lazy_features(g, feature_names):
"""Assign lazy features to the ``srcdata`` of the input graph for prefetching optimization.
When used in a :class:`~dgl.dataloading.Sampler`, lazy features mark which data
should be fetched before computation in model. See :ref:`guide-minibatch-prefetching`
for a detailed explanation.
If the graph is homogeneous, this is equivalent to:
.. code:: python
g.srcdata.update({k: LazyFeature(k, g.srcdata[dgl.NID]) for k in feature_names})
If the graph is heterogeneous, this is equivalent to:
.. code:: python
for type_, names in feature_names.items():
g.srcnodes[type_].data.update(
{k: LazyFeature(k, g.srcnodes[type_].data[dgl.NID]) for k in names})
Parameters
----------
g : DGLGraph
The graph.
feature_names : list[str] or dict[str, list[str]]
The feature names to prefetch.
See also
--------
dgl.LazyFeature
"""
return _set_lazy_features(g.srcnodes, g.srcdata, feature_names)
def set_dst_lazy_features(g, feature_names):
"""Assign lazy features to the ``dstdata`` of the input graph for prefetching optimization.
When used in a :class:`~dgl.dataloading.Sampler`, lazy features mark which data
should be fetched before computation in model. See :ref:`guide-minibatch-prefetching`
for a detailed explanation.
If the graph is homogeneous, this is equivalent to:
.. code:: python
g.dstdata.update({k: LazyFeature(k, g.dstdata[dgl.NID]) for k in feature_names})
If the graph is heterogeneous, this is equivalent to:
.. code:: python
for type_, names in feature_names.items():
g.dstnodes[type_].data.update(
{k: LazyFeature(k, g.dstnodes[type_].data[dgl.NID]) for k in names})
Parameters
----------
g : DGLGraph
The graph.
feature_names : list[str] or dict[str, list[str]]
The feature names to prefetch.
See also
--------
dgl.LazyFeature
"""
return _set_lazy_features(g.dstnodes, g.dstdata, feature_names)
class Sampler(object):
"""Base class for graph samplers.
All graph samplers must subclass this class and override the ``sample``
method.
.. code:: python
from dgl.dataloading import Sampler
class SubgraphSampler(Sampler):
def __init__(self):
super().__init__()
def sample(self, g, indices):
return g.subgraph(indices)
"""
def sample(self, g, indices):
"""Abstract sample method.
Parameters
----------
g : DGLGraph
The graph.
indices : object
Any object representing the indices selected in the current minibatch.
"""
raise NotImplementedError
class BlockSampler(Sampler):
"""Base class for sampling mini-batches in the form of Message-passing
Flow Graphs (MFGs).
It provides prefetching options to fetch the node features for the first MFG's ``srcdata``,
the node labels for the last MFG's ``dstdata`` and the edge features of all MFG's ``edata``.
Parameters
----------
prefetch_node_feats : list[str] or dict[str, list[str]], optional
The node data to prefetch for the first MFG.
DGL will populate the first layer's MFG's ``srcnodes`` and ``srcdata`` with
the node data of the given names from the original graph.
prefetch_labels : list[str] or dict[str, list[str]], optional
The node data to prefetch for the last MFG.
DGL will populate the last layer's MFG's ``dstnodes`` and ``dstdata`` with
the node data of the given names from the original graph.
prefetch_edge_feats : list[str] or dict[etype, list[str]], optional
The edge data names to prefetch for all the MFGs.
DGL will populate every MFG's ``edges`` and ``edata`` with the edge data
of the given names from the original graph.
output_device : device, optional
The device of the output subgraphs or MFGs. Default is the same as the
minibatch of seed nodes.
"""
def __init__(
self,
prefetch_node_feats=None,
prefetch_labels=None,
prefetch_edge_feats=None,
output_device=None,
):
super().__init__()
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."""
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, exclude_eids=None
): # pylint: disable=arguments-differ
"""Sample a list of blocks from the given seed nodes."""
result = self.sample_blocks(g, seed_nodes, exclude_eids=exclude_eids)
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()
}
for k, v in reverse_etype_map.items():
if k in exclude_eids:
if v in exclude_eids:
exclude_eids[v] = F.unique(
F.cat((exclude_eids[k], exclude_eids[v]), dim=0)
)
else:
exclude_eids[v] = exclude_eids[k]
return exclude_eids
def _find_exclude_eids(g, exclude_mode, eids, **kwargs):
if exclude_mode is None:
return None
elif callable(exclude_mode):
return exclude_mode(eids)
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 :
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]``.
callable
Any function that takes in a single argument :attr:`seed_edges` and returns
a tensor or dict of tensors.
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 and output_device is not None:
exclude_eids = recursive_apply(
exclude_eids, lambda x: F.copy_to(x, output_device)
)
return exclude_eids
class EdgePredictionSampler(Sampler):
"""Sampler class that wraps an existing sampler for node classification into another
one for edge classification or link prediction.
See also
--------
as_edge_prediction_sampler
"""
def __init__(
self,
sampler,
exclude=None,
reverse_eids=None,
reverse_etypes=None,
negative_sampler=None,
prefetch_labels=None,
):
super().__init__()
# Check if the sampler's sample method has an optional third argument.
argspec = inspect.getfullargspec(sampler.sample)
if len(argspec.args) < 4: # ['self', 'g', 'indices', 'exclude_eids']
raise TypeError(
"This sampler does not support edge or link prediction; please add an"
"optional third argument for edge IDs to exclude in its sample() method."
)
self.reverse_eids = reverse_eids
self.reverse_etypes = reverse_etypes
self.exclude = exclude
self.sampler = sampler
self.negative_sampler = negative_sampler
self.prefetch_labels = prefetch_labels or []
self.output_device = 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])
ctx = context_of(seed_edges) if seed_edges is not None else g.device
neg_edges = {
etype: neg_srcdst.get(
etype,
(
F.copy_to(F.tensor([], dtype), ctx=ctx),
F.copy_to(F.tensor([], dtype), ctx=ctx),
),
)
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]
set_edge_lazy_features(pair_graph, self.prefetch_labels)
# In-place updates
return result
def sample(self, g, seed_edges): # pylint: disable=arguments-differ
"""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.
"""
if isinstance(seed_edges, Mapping):
seed_edges = {
g.to_canonical_etype(k): v for k, v in seed_edges.items()
}
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.sampler.sample(
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)
)
def as_edge_prediction_sampler(
sampler,
exclude=None,
reverse_eids=None,
reverse_etypes=None,
negative_sampler=None,
prefetch_labels=None,
):
"""Create an edge-wise sampler from a node-wise sampler.
For each batch of edges, the sampler applies the provided node-wise sampler to
their source and destination nodes to extract subgraphs. It also generates negative
edges if a negative sampler is provided, and extract subgraphs for their incident
nodes as well.
For each iteration, the sampler 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.
* The subgraphs or MFGs returned by the provided node-wise sampler, generated
from the incident nodes of the edges in the minibatch (as well as those of the
negative edges if applicable).
Parameters
----------
sampler : Sampler
The node-wise sampler object. It additionally requires that the :attr:`sample`
method must have an optional third argument :attr:`exclude_eids` representing the
edge IDs to exclude from neighborhood. The argument will be either a tensor
for homogeneous graphs or a dict of edge types and tensors for heterogeneous
graphs.
exclude : Union[str, callable], optional
Whether and how to exclude dependencies related to the sampled edges in the
minibatch. Possible values are
* None, for not excluding any edges.
* ``self``, for excluding the edges in the current minibatch.
* ``reverse_id``, for excluding not only the edges in the current minibatch but
also their reverse edges according to the ID mapping in the argument
:attr:`reverse_eids`.
* ``reverse_types``, for excluding not only the edges in the current minibatch
but also their reverse edges stored in another type according to
the argument :attr:`reverse_etypes`.
* User-defined exclusion rule. It is a callable with edges in the current
minibatch as a single argument and should return the edges to be excluded.
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.
reverse_etypes : dict[etype, etype], optional
The mapping from the original edge types to their reverse edge types.
negative_sampler : callable, optional
The negative sampler.
prefetch_labels : list[str] or dict[etype, list[str]], optional
The edge labels to prefetch for the returned positive pair graph.
See :ref:`guide-minibatch-prefetching` for a detailed explanation of prefetching.
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.
Given an array of source node IDs ``src`` and another array of destination
node IDs ``dst``, the following code creates a bidirectional graph:
>>> g = dgl.graph((torch.cat([src, dst]), torch.cat([dst, src])))
Edge :math:`i`'s reverse edge in the graph above is edge :math:`i + |E|`. Therefore, we can
create a reverse edge mapping ``reverse_eids`` by:
>>> E = len(src)
>>> reverse_eids = torch.cat([torch.arange(E, 2 * E), torch.arange(0, E)])
By passing ``reverse_eids`` to the edge sampler, the edges in the current mini-batch and their
reversed edges will be excluded from the extracted subgraphs to avoid information leakage.
>>> sampler = dgl.dataloading.as_edge_prediction_sampler(
... dgl.dataloading.NeighborSampler([15, 10, 5]),
... exclude='reverse_id', reverse_eids=reverse_eids)
>>> dataloader = dgl.dataloading.DataLoader(
... g, train_eid, sampler,
... 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)
For link prediction, one can provide a negative sampler to sample negative edges.
The code below uses DGL's :class:`~dgl.dataloading.negative_sampler.Uniform`
to generate 5 negative samples per edge:
>>> neg_sampler = dgl.dataloading.negative_sampler.Uniform(5)
>>> sampler = dgl.dataloading.as_edge_prediction_sampler(
... dgl.dataloading.NeighborSampler([15, 10, 5]),
... sampler, exclude='reverse_id', reverse_eids=reverse_eids,
... negative_sampler=neg_sampler)
>>> dataloader = dgl.dataloading.DataLoader(
... g, train_eid, sampler,
... 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)
For heterogeneous graphs, reverse edges may belong to a different relation. For example,
the relations "user-click-item" and "item-click-by-user" in the graph below are
mutual reverse.
>>> g = dgl.heterograph({
... ('user', 'click', 'item'): (user, item),
... ('item', 'clicked-by', 'user'): (item, user)})
To correctly exclude edges from each mini-batch, set ``exclude='reverse_types'`` and
pass a dictionary ``{'click': 'clicked-by', 'clicked-by': 'click'}`` to the
``reverse_etypes`` argument.
>>> sampler = dgl.dataloading.as_edge_prediction_sampler(
... dgl.dataloading.NeighborSampler([15, 10, 5]),
... exclude='reverse_types',
... reverse_etypes={'click': 'clicked-by', 'clicked-by': 'click'})
>>> dataloader = dgl.dataloading.DataLoader(
... g, {'click': train_eid}, sampler,
... 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)
For link prediction, provide a negative sampler to generate negative samples:
>>> neg_sampler = dgl.dataloading.negative_sampler.Uniform(5)
>>> sampler = dgl.dataloading.as_edge_prediction_sampler(
... dgl.dataloading.NeighborSampler([15, 10, 5]),
... exclude='reverse_types',
... reverse_etypes={'click': 'clicked-by', 'clicked-by': 'click'},
... negative_sampler=neg_sampler)
>>> dataloader = dgl.dataloading.DataLoader(
... g, train_eid, sampler,
... 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)
"""
return EdgePredictionSampler(
sampler,
exclude=exclude,
reverse_eids=reverse_eids,
reverse_etypes=reverse_etypes,
negative_sampler=negative_sampler,
prefetch_labels=prefetch_labels,
)
python/dgl/dataloading/capped_neighbor_sampler.py
+190
View File
@@ -0,0 +1,190 @@
"""Capped neighbor sampler."""
from collections import defaultdict
import numpy as np
import torch
from ..sampling.utils import EidExcluder
from .base import Sampler, set_edge_lazy_features, set_node_lazy_features
class CappedNeighborSampler(Sampler):
"""Subgraph sampler that sets an upper bound on the number of nodes included in
each layer of the sampled subgraph. At each layer, the frontier is randomly
subsampled. Rare node types can also be upsampled by taking the scaled square
root of the sampling probabilities. The sampler returns the subgraph induced by
all the sampled nodes.
This code was contributed by a community member
([@ayushnoori](https://github.com/ayushnoori)). There aren't currently any unit
tests in place to verify its functionality, so please be cautious if you need
to make any changes to the code's logic.
Parameters
----------
fanouts : list[int] or 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.
fixed_k : int
The number of nodes to sample for each GNN layer.
upsample_rare_types : bool
Whether or not to upsample rare node types.
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.
"""
def __init__(
self,
fanouts,
fixed_k,
upsample_rare_types,
replace=False,
prob=None,
prefetch_node_feats=None,
prefetch_edge_feats=None,
output_device=None,
):
super().__init__()
self.fanouts = fanouts
self.replace = replace
self.fixed_k = fixed_k
self.upsample_rare_types = upsample_rare_types
self.prob = prob
self.prefetch_node_feats = prefetch_node_feats
self.prefetch_edge_feats = prefetch_edge_feats
self.output_device = output_device
def sample(
self, g, indices, exclude_eids=None
): # pylint: disable=arguments-differ
"""Sampling function.
Parameters
----------
g : DGLGraph
The graph to sample from.
indices : Tensor or dict[str, Tensor]
Nodes which induce the subgraph.
exclude_eids : Tensor or dict[etype, Tensor], optional
The edges to exclude from the sampled subgraph.
Returns
-------
input_nodes : Tensor or dict[str, Tensor]
The node IDs inducing the subgraph.
output_nodes : Tensor or dict[str, Tensor]
The node IDs that are sampled in this minibatch.
subg : DGLGraph
The subgraph itself.
"""
# Define empty dictionary to store reached nodes.
output_nodes = indices
all_reached_nodes = [indices]
# Iterate over fanout.
for fanout in reversed(self.fanouts):
# Sample frontier.
frontier = g.sample_neighbors(
indices,
fanout,
output_device=self.output_device,
replace=self.replace,
prob=self.prob,
exclude_edges=exclude_eids,
)
# Get reached nodes.
curr_reached = defaultdict(list)
for c_etype in frontier.canonical_etypes:
(src_type, _, _) = c_etype
src, _ = frontier.edges(etype=c_etype)
curr_reached[src_type].append(src)
# De-duplication.
curr_reached = {
ntype: torch.unique(torch.cat(srcs))
for ntype, srcs in curr_reached.items()
}
# Generate type sampling probabilties.
type_count = {
node_type: indices.shape[0]
for node_type, indices in curr_reached.items()
}
total_count = sum(type_count.values())
probs = {
node_type: count / total_count
for node_type, count in type_count.items()
}
# Upsample rare node types.
if self.upsample_rare_types:
# Take scaled square root of probabilities.
prob_dist = list(probs.values())
prob_dist = np.sqrt(prob_dist)
prob_dist = prob_dist / prob_dist.sum()
# Update probabilities.
probs = {
node_type: prob_dist[i]
for i, node_type in enumerate(probs.keys())
}
# Generate node counts per type.
n_per_type = {
node_type: int(self.fixed_k * prob)
for node_type, prob in probs.items()
}
remainder = self.fixed_k - sum(n_per_type.values())
for _ in range(remainder):
node_type = np.random.choice(
list(probs.keys()), p=list(probs.values())
)
n_per_type[node_type] += 1
# Downsample nodes.
curr_reached_k = {}
for node_type, node_ids in curr_reached.items():
# Get number of total nodes and number to sample.
num_nodes = node_ids.shape[0]
n_to_sample = min(num_nodes, n_per_type[node_type])
# Downsample nodes of current type.
random_indices = torch.randperm(num_nodes)[:n_to_sample]
curr_reached_k[node_type] = node_ids[random_indices]
# Update seed nodes.
indices = curr_reached_k
all_reached_nodes.append(curr_reached_k)
# Merge all reached nodes before sending to `DGLGraph.subgraph`.
merged_nodes = {}
for ntype in g.ntypes:
merged_nodes[ntype] = torch.unique(
torch.cat(
[reached.get(ntype, []) for reached in all_reached_nodes]
)
)
subg = g.subgraph(
merged_nodes, relabel_nodes=True, output_device=self.output_device
)
if exclude_eids is not None:
subg = EidExcluder(exclude_eids)(subg)
set_node_lazy_features(subg, self.prefetch_node_feats)
set_edge_lazy_features(subg, self.prefetch_edge_feats)
return indices, output_nodes, subg
python/dgl/dataloading/cluster_gcn.py
+155
View File
@@ -0,0 +1,155 @@
"""Cluster-GCN samplers."""
import os
import pickle
import numpy as np
from .. import backend as F
from ..base import DGLError
from ..partition import metis_partition_assignment
from .base import Sampler, set_edge_lazy_features, set_node_lazy_features
class ClusterGCNSampler(Sampler):
"""Cluster sampler from `Cluster-GCN: An Efficient Algorithm for Training
Deep and Large Graph Convolutional Networks
<https://arxiv.org/abs/1905.07953>`__
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.
The sampler then selects the graph partitions according to the provided
partition IDs, take the union of all nodes in those partitions, and return an
induced subgraph in its :attr:`sample` method.
Parameters
----------
g : DGLGraph
The original graph. Must be homogeneous and on CPU.
k : int
The number of partitions.
cache_path : str
The path to the cache directory for storing the partition result.
balance_ntypes, balkance_edges, mode :
Passed to :func:`dgl.metis_partition_assignment`.
prefetch_ndata : list[str], optional
The node data to prefetch for the subgraph.
See :ref:`guide-minibatch-prefetching` for a detailed explanation of prefetching.
prefetch_edata : list[str], optional
The edge data to prefetch for the subgraph.
See :ref:`guide-minibatch-prefetching` for a detailed explanation of prefetching.
output_device : device, optional
The device of the output subgraphs or MFGs. Default is the same as the
minibatch of partition indices.
Examples
--------
**Node classification**
With this sampler, the data loader will accept the list of partition IDs as
indices to iterate over. For instance, the following code first splits the
graph into 1000 partitions using METIS, and at each iteration it gets a subgraph
induced by the nodes covered by 20 randomly selected partitions.
>>> num_parts = 1000
>>> sampler = dgl.dataloading.ClusterGCNSampler(g, num_parts)
>>> dataloader = dgl.dataloading.DataLoader(
... g, torch.arange(num_parts), sampler,
... batch_size=20, shuffle=True, drop_last=False, num_workers=4)
>>> for subg in dataloader:
... train_on(subg)
"""
def __init__(
self,
g,
k,
cache_path="cluster_gcn.pkl",
balance_ntypes=None,
balance_edges=False,
mode="k-way",
prefetch_ndata=None,
prefetch_edata=None,
output_device=None,
):
super().__init__()
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_node_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_ndata = prefetch_ndata or []
self.prefetch_edata = prefetch_edata or []
self.output_device = output_device
def sample(self, g, partition_ids): # pylint: disable=arguments-differ
"""Sampling function.
Parameters
----------
g : DGLGraph
The graph to sample from.
partition_ids : Tensor
A 1-D integer tensor of partition IDs.
Returns
-------
DGLGraph
The sampled subgraph.
"""
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_ndata)
set_edge_lazy_features(sg, self.prefetch_edata)
return sg
python/dgl/dataloading/dataloader.py
+1560
View File
File diff suppressed because it is too large Load Diff
python/dgl/dataloading/graphsaint.py
+163
View File
@@ -0,0 +1,163 @@
"""GraphSAINT samplers."""
from ..base import DGLError
from ..random import choice
from ..sampling import pack_traces, random_walk
from .base import Sampler, set_edge_lazy_features, set_node_lazy_features
try:
import torch
except ImportError:
pass
class SAINTSampler(Sampler):
"""Random node/edge/walk sampler from
`GraphSAINT: Graph Sampling Based Inductive Learning Method
<https://arxiv.org/abs/1907.04931>`__
For each call, the sampler samples a node subset and then returns a node induced subgraph.
There are three options for sampling node subsets:
- For :attr:`'node'` sampler, the probability to sample a node is in proportion
to its out-degree.
- The :attr:`'edge'` sampler first samples an edge subset and then use the
end nodes of the edges.
- The :attr:`'walk'` sampler uses the nodes visited by random walks. It uniformly selects
a number of root nodes and then performs a fixed-length random walk from each root node.
Parameters
----------
mode : str
The sampler to use, which can be :attr:`'node'`, :attr:`'edge'`, or :attr:`'walk'`.
budget : int or tuple[int]
Sampler configuration.
- For :attr:`'node'` sampler, budget specifies the number of nodes
in each sampled subgraph.
- For :attr:`'edge'` sampler, budget specifies the number of edges
to sample for inducing a subgraph.
- For :attr:`'walk'` sampler, budget is a tuple. budget[0] specifies
the number of root nodes to generate random walks. budget[1] specifies
the length of a random walk.
cache : bool, optional
If False, it will not cache the probability arrays for sampling. Setting
it to False is required if you want to use the sampler across different graphs.
prefetch_ndata : list[str], optional
The node data to prefetch for the subgraph.
See :ref:`guide-minibatch-prefetching` for a detailed explanation of prefetching.
prefetch_edata : list[str], optional
The edge data to prefetch for the subgraph.
See :ref:`guide-minibatch-prefetching` for a detailed explanation of prefetching.
output_device : device, optional
The device of the output subgraphs.
Examples
--------
>>> import torch
>>> from dgl.dataloading import SAINTSampler, DataLoader
>>> num_iters = 1000
>>> sampler = SAINTSampler(mode='node', budget=6000)
>>> # Assume g.ndata['feat'] and g.ndata['label'] hold node features and labels
>>> dataloader = DataLoader(g, torch.arange(num_iters), sampler, num_workers=4)
>>> for subg in dataloader:
... train_on(subg)
"""
def __init__(
self,
mode,
budget,
cache=True,
prefetch_ndata=None,
prefetch_edata=None,
output_device="cpu",
):
super().__init__()
self.budget = budget
if mode == "node":
self.sampler = self.node_sampler
elif mode == "edge":
self.sampler = self.edge_sampler
elif mode == "walk":
self.sampler = self.walk_sampler
else:
raise DGLError(
f"Expect mode to be 'node', 'edge' or 'walk', got {mode}."
)
self.cache = cache
self.prob = None
self.prefetch_ndata = prefetch_ndata or []
self.prefetch_edata = prefetch_edata or []
self.output_device = output_device
def node_sampler(self, g):
"""Node ID sampler for random node sampler"""
# Alternatively, this can be realized by uniformly sampling an edge subset,
# and then take the src node of the sampled edges. However, the number of edges
# is typically much larger than the number of nodes.
if self.cache and self.prob is not None:
prob = self.prob
else:
prob = g.out_degrees().float().clamp(min=1)
if self.cache:
self.prob = prob
return (
torch.multinomial(prob, num_samples=self.budget, replacement=True)
.unique()
.type(g.idtype)
)
def edge_sampler(self, g):
"""Node ID sampler for random edge sampler"""
src, dst = g.edges()
if self.cache and self.prob is not None:
prob = self.prob
else:
in_deg = g.in_degrees().float().clamp(min=1)
out_deg = g.out_degrees().float().clamp(min=1)
# We can reduce the sample space by half if graphs are always symmetric.
prob = 1.0 / in_deg[dst.long()] + 1.0 / out_deg[src.long()]
prob /= prob.sum()
if self.cache:
self.prob = prob
sampled_edges = torch.unique(
choice(len(prob), size=self.budget, prob=prob)
)
sampled_nodes = torch.cat([src[sampled_edges], dst[sampled_edges]])
return sampled_nodes.unique().type(g.idtype)
def walk_sampler(self, g):
"""Node ID sampler for random walk sampler"""
num_roots, walk_length = self.budget
sampled_roots = torch.randint(0, g.num_nodes(), (num_roots,))
traces, types = random_walk(g, nodes=sampled_roots, length=walk_length)
sampled_nodes, _, _, _ = pack_traces(traces, types)
return sampled_nodes.unique().type(g.idtype)
def sample(self, g, indices):
"""Sampling function
Parameters
----------
g : DGLGraph
The graph to sample from.
indices : Tensor
Placeholder not used.
Returns
-------
DGLGraph
The sampled subgraph.
"""
node_ids = self.sampler(g)
sg = g.subgraph(
node_ids, relabel_nodes=True, output_device=self.output_device
)
set_node_lazy_features(sg, self.prefetch_ndata)
set_edge_lazy_features(sg, self.prefetch_edata)
return sg
python/dgl/dataloading/labor_sampler.py
+255
View File
@@ -0,0 +1,255 @@
#
# Copyright (c) 2022 by Contributors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Based off of neighbor_sampler.py
#
"""Data loading components for labor sampling"""
from numpy.random import default_rng
from .. import backend as F
from ..base import EID, NID
from ..random import choice
from ..transforms import to_block
from .base import BlockSampler
class LaborSampler(BlockSampler):
"""Sampler that builds computational dependency of node representations via
labor sampling for multilayer GNN from the NeurIPS 2023 paper
`Layer-Neighbor Sampling -- Defusing Neighborhood Explosion in GNNs
<https://arxiv.org/abs/2210.13339>`__
This sampler will make every node gather messages from a fixed number of
neighbors per edge type. The neighbors are picked uniformly with default
parameters. For every vertex t that will be considered to be sampled, there
will be a single random variate r_t.
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.
edge_dir : str, default ``'in'``
Can be either ``'in'`` where the neighbors will be sampled according to
incoming edges, or ``'out'`` otherwise, same as
:func:`dgl.sampling.sample_neighbors`.
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. In this case, the returned
blocks edata include ``'edge_weights'`` that needs to be used in the
message passing operation.
importance_sampling : int, default ``0``
Whether to use importance sampling or uniform sampling, use of negative
values optimizes importance sampling probabilities until convergence
while use of positive values runs optimization steps that many times.
If the value is i, then LABOR-i variant is used. When used with a
nonzero parameter, the returned blocks edata include ``'edge_weights'``
that needs to be used in the message passing operation.
layer_dependency : bool, default ``False``
Specifies whether different layers should use same random variates.
Results into a reduction in the number of vertices sampled, but may
degrade the quality slightly.
batch_dependency : int, default ``1``
Specifies whether different minibatches should use similar random
variates. Results in a higher temporal access locality of sampled
vertices, but may degrade the quality slightly.
prefetch_node_feats : list[str] or dict[ntype, list[str]], optional
The source node data to prefetch for the first MFG, corresponding to the
input node features necessary for the first GNN layer.
prefetch_labels : list[str] or dict[ntype, list[str]], optional
The destination node data to prefetch for the last MFG, corresponding to
the node labels of the minibatch.
prefetch_edge_feats : list[str] or dict[etype, list[str]], optional
The edge data names to prefetch for all the MFGs, corresponding to the
edge features necessary for all GNN layers.
output_device : device, optional
The device of the output subgraphs or MFGs. Default is the same as the
minibatch of seed nodes.
Examples
--------
**Node classification**
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.LaborSampler([5, 10, 15])
>>> dataloader = dgl.dataloading.DataLoader(
... 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.LaborSampler([
... {('user', 'follows', 'user'): 5,
... ('user', 'plays', 'game'): 4,
... ('game', 'played-by', 'user'): 3}] * 3)
If you would like non-uniform labor sampling:
>>> # any non-negative 1D vector works
>>> g.edata['p'] = torch.rand(g.num_edges())
>>> sampler = dgl.dataloading.LaborSampler([5, 10, 15], prob='p')
**Edge classification and link prediction**
This class can also work for edge classification and link prediction
together with :func:`as_edge_prediction_sampler`.
>>> sampler = dgl.dataloading.LaborSampler([5, 10, 15])
>>> sampler = dgl.dataloading.as_edge_prediction_sampler(sampler)
>>> dataloader = dgl.dataloading.DataLoader(
... g, train_eid, sampler,
... batch_size=1024, shuffle=True, drop_last=False, num_workers=4)
See the documentation :func:`as_edge_prediction_sampler` for more details.
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,
importance_sampling=0,
layer_dependency=False,
batch_dependency=1,
prefetch_node_feats=None,
prefetch_labels=None,
prefetch_edge_feats=None,
output_device=None,
):
super().__init__(
prefetch_node_feats=prefetch_node_feats,
prefetch_labels=prefetch_labels,
prefetch_edge_feats=prefetch_edge_feats,
output_device=output_device,
)
self.fanouts = fanouts
self.edge_dir = edge_dir
self.prob = prob
self.importance_sampling = importance_sampling
self.layer_dependency = layer_dependency
self.cnt = F.zeros(2, F.int64, F.cpu())
self.cnt[0] = -1
self.cnt[1] = batch_dependency
self.random_seed = F.zeros(
2 if self.cnt[1] > 1 else 1, F.int64, F.cpu()
)
self.set_seed(None if batch_dependency > 0 else choice(1e18, 1).item())
def set_seed(self, random_seed=None):
"""Updates the underlying seed for the sampler
Calling this function enforces the sampling algorithm to use the same
seed on every edge type. This can reduce the number of nodes being
sampled because the passed random_seed makes it so that for any seed
vertex ``s`` and its neighbor ``t``, the rolled random variate ``r_t``
is the same for any instance of this class with the same random seed.
When sampling as part of the same batch, one would want identical seeds
so that LABOR can globally sample. One example is that for heterogenous
graphs, there is a single random seed passed for each edge type. This
will sample much fewer vertices compared to having unique random seeds
for each edge type. If one called this function individually for each
edge type for a heterogenous graph with different random seeds, then it
would run LABOR locally for each edge type, resulting into a larger
number of vertices being sampled.
If this function is called without any parameters, we get the random
seed by getting a random number from DGL. Call this function if multiple
instances of LaborSampler are used to sample as part of a single batch.
Parameters
----------
random_seed : int, default ``None``
The random seed to be used for next sampling call.
"""
if random_seed is None:
self.cnt[0] += 1
if self.cnt[1] > 0 and self.cnt[0] % self.cnt[1] == 0:
if self.cnt[0] <= 0 or self.cnt[1] <= 1:
if not hasattr(self, "rng"):
self.rng = default_rng(choice(1e18, 1).item())
self.random_seed[0] = self.rng.integers(1e18)
if self.cnt[1] > 1:
self.random_seed[1] = self.rng.integers(1e18)
else:
self.random_seed[0] = self.random_seed[1]
self.random_seed[1] = self.rng.integers(1e18)
else:
self.rng = default_rng(random_seed)
self.random_seed[0] = self.rng.integers(1e18)
if self.cnt[1] > 1:
self.random_seed[1] = self.rng.integers(1e18)
self.cnt[0] = 0
def sample_blocks(self, g, seed_nodes, exclude_eids=None):
output_nodes = seed_nodes
blocks = []
for i, fanout in enumerate(reversed(self.fanouts)):
random_seed_i = F.zerocopy_to_dgl_ndarray(
self.random_seed + (i if not self.layer_dependency else 0)
)
if self.cnt[1] <= 1:
seed2_contr = 0
else:
seed2_contr = ((self.cnt[0] % self.cnt[1]) / self.cnt[1]).item()
frontier, importances = g.sample_labors(
seed_nodes,
fanout,
edge_dir=self.edge_dir,
prob=self.prob,
importance_sampling=self.importance_sampling,
random_seed=random_seed_i,
seed2_contribution=seed2_contr,
output_device=self.output_device,
exclude_edges=exclude_eids,
)
eid = frontier.edata[EID]
block = to_block(
frontier, seed_nodes, include_dst_in_src=True, src_nodes=None
)
block.edata[EID] = eid
if len(g.canonical_etypes) > 1:
for etype, importance in zip(g.canonical_etypes, importances):
if importance.shape[0] == block.num_edges(etype):
block.edata["edge_weights"][etype] = importance
elif importances[0].shape[0] == block.num_edges():
block.edata["edge_weights"] = importances[0]
seed_nodes = block.srcdata[NID]
blocks.insert(0, block)
self.set_seed()
return seed_nodes, output_nodes, blocks
python/dgl/dataloading/negative_sampler.py
+126
View File
@@ -0,0 +1,126 @@
"""Negative samplers"""
from collections.abc import Mapping
from .. import backend as F
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.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])
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.num_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: False)
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):
self.k = k
self.exclude_self_loops = exclude_self_loops
self.replace = replace
def _generate(self, g, eids, canonical_etype):
return g.global_uniform_negative_sampling(
len(eids) * self.k,
self.exclude_self_loops,
self.replace,
canonical_etype,
)
python/dgl/dataloading/neighbor_sampler.py
+246
View File
@@ -0,0 +1,246 @@
"""Data loading components for neighbor sampling"""
from .. import backend as F
from ..base import EID, NID
from ..heterograph import DGLGraph
from ..transforms import to_block
from ..utils import get_num_threads
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.
edge_dir : str, default ``'in'``
Can be either ``'in' `` where the neighbors will be sampled according to
incoming edges, or ``'out'`` otherwise, same as :func:`dgl.sampling.sample_neighbors`.
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.
This argument is mutually exclusive with :attr:`mask`. If you want to
specify both a mask and a probability, consider multiplying the probability
with the mask instead.
mask : str, optional
If given, a neighbor could be picked only if the edge mask with the given
name in ``g.edata`` is True. The data must be boolean on each edge.
This argument is mutually exclusive with :attr:`prob`. If you want to
specify both a mask and a probability, consider multiplying the probability
with the mask instead.
replace : bool, default False
Whether to sample with replacement
prefetch_node_feats : list[str] or dict[ntype, list[str]], optional
The source node data to prefetch for the first MFG, corresponding to the
input node features necessary for the first GNN layer.
prefetch_labels : list[str] or dict[ntype, list[str]], optional
The destination node data to prefetch for the last MFG, corresponding to
the node labels of the minibatch.
prefetch_edge_feats : list[str] or dict[etype, list[str]], optional
The edge data names to prefetch for all the MFGs, corresponding to the
edge features necessary for all GNN layers.
output_device : device, optional
The device of the output subgraphs or MFGs. Default is the same as the
minibatch of seed nodes.
fused : bool, default True
If True and device is CPU fused sample neighbors is invoked. This version
requires seed_nodes to be unique
Examples
--------
**Node classification**
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.DataLoader(
... 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')
Or sampling on edge masks:
>>> g.edata['mask'] = torch.rand(g.num_edges()) < 0.2 # any 1D boolean mask works
>>> sampler = dgl.dataloading.NeighborSampler([5, 10, 15], prob='mask')
**Edge classification and link prediction**
This class can also work for edge classification and link prediction together
with :func:`as_edge_prediction_sampler`.
>>> sampler = dgl.dataloading.NeighborSampler([5, 10, 15])
>>> sampler = dgl.dataloading.as_edge_prediction_sampler(sampler)
>>> dataloader = dgl.dataloading.DataLoader(
... g, train_eid, sampler,
... batch_size=1024, shuffle=True, drop_last=False, num_workers=4)
See the documentation :func:`as_edge_prediction_sampler` for more details.
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,
mask=None,
replace=False,
prefetch_node_feats=None,
prefetch_labels=None,
prefetch_edge_feats=None,
output_device=None,
fused=True,
):
super().__init__(
prefetch_node_feats=prefetch_node_feats,
prefetch_labels=prefetch_labels,
prefetch_edge_feats=prefetch_edge_feats,
output_device=output_device,
)
self.fanouts = fanouts
self.edge_dir = edge_dir
if mask is not None and prob is not None:
raise ValueError(
"Mask and probability arguments are mutually exclusive. "
"Consider multiplying the probability with the mask "
"to achieve the same goal."
)
self.prob = prob or mask
self.replace = replace
self.fused = fused
self.mapping = {}
self.g = None
def sample_blocks(self, g, seed_nodes, exclude_eids=None):
output_nodes = seed_nodes
blocks = []
# sample_neighbors_fused function requires multithreading to be more efficient
# than sample_neighbors
if self.fused and get_num_threads() > 1:
cpu = F.device_type(g.device) == "cpu"
if isinstance(seed_nodes, dict):
for ntype in list(seed_nodes.keys()):
if not cpu:
break
cpu = (
cpu and F.device_type(seed_nodes[ntype].device) == "cpu"
)
else:
cpu = cpu and F.device_type(seed_nodes.device) == "cpu"
if cpu and isinstance(g, DGLGraph) and F.backend_name == "pytorch":
if self.g != g:
self.mapping = {}
self.g = g
for fanout in reversed(self.fanouts):
block = g.sample_neighbors_fused(
seed_nodes,
fanout,
edge_dir=self.edge_dir,
prob=self.prob,
replace=self.replace,
exclude_edges=exclude_eids,
mapping=self.mapping,
)
seed_nodes = block.srcdata[NID]
blocks.insert(0, block)
return seed_nodes, output_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,
)
block = to_block(frontier, seed_nodes)
# If sampled from graphbolt-backed DistGraph, `EID` may not be in
# the block. If not exists, we should remove it from the block.
if EID in frontier.edata.keys():
block.edata[EID] = frontier.edata[EID]
else:
del block.edata[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
----------
num_layers : int
The number of GNN layers to sample.
kwargs :
Passed to :class:`dgl.dataloading.NeighborSampler`.
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.DataLoader(
... 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, **kwargs):
super().__init__([-1] * num_layers, **kwargs)
python/dgl/dataloading/shadow.py
+132
View File
@@ -0,0 +1,132 @@
"""ShaDow-GNN subgraph samplers."""
from .. import transforms
from ..base import NID
from ..sampling.utils import EidExcluder
from .base import Sampler, set_edge_lazy_features, set_node_lazy_features
class ShaDowKHopSampler(Sampler):
"""K-hop subgraph sampler from `Deep Graph Neural Networks with Shallow
Subgraph Samplers <https://arxiv.org/abs/2012.01380>`__.
It performs node-wise neighbor sampling and returns the 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.
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
--------
**Node classification**
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.DataLoader(
... 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.ShaDowKHopSampler([5, 10, 15], prob='p')
"""
def __init__(
self,
fanouts,
replace=False,
prob=None,
prefetch_node_feats=None,
prefetch_edge_feats=None,
output_device=None,
):
super().__init__()
self.fanouts = fanouts
self.replace = replace
self.prob = prob
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
): # pylint: disable=arguments-differ
"""Sampling function.
Parameters
----------
g : DGLGraph
The graph to sample nodes from.
seed_nodes : Tensor or dict[str, Tensor]
The nodes sampled in the current minibatch.
exclude_eids : Tensor or dict[etype, Tensor], optional
The edges to exclude from neighborhood expansion.
Returns
-------
input_nodes, output_nodes, subg
A triplet containing (1) the node IDs inducing the subgraph, (2) the node
IDs that are sampled in this minibatch, and (3) the subgraph itself.
"""
output_nodes = seed_nodes
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_eids,
)
block = transforms.to_block(frontier, seed_nodes)
seed_nodes = block.srcdata[NID]
subg = g.subgraph(
seed_nodes, relabel_nodes=True, output_device=self.output_device
)
if exclude_eids is not None:
subg = EidExcluder(exclude_eids)(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
python/dgl/dataloading/spot_target.py
+89
View File
@@ -0,0 +1,89 @@
"""SpotTarget: Target edge excluder for link prediction"""
import torch
from .base import find_exclude_eids
class SpotTarget(object):
"""Callable excluder object to exclude the edges by the degree threshold.
Besides excluding all the edges or given edges in the edge sampler
``dgl.dataloading.as_edge_prediction_sampler`` in link prediction training,
this excluder can extend the exclusion function by only excluding the edges incident
to low-degree nodes in the graph to bring the performance increase in training
link prediction model. This function will exclude the edge if incident to at least
one node with degree larger or equal to ``degree_threshold``. The performance
boost by excluding the target edges incident to low-degree nodes can be found
in this paper: https://arxiv.org/abs/2306.00899
Parameters
----------
g : DGLGraph
The graph.
exclude : Union[str, callable]
Whether and how to exclude dependencies related to the sampled edges in the
minibatch. Possible values are
* ``self``, for excluding the edges in the current minibatch.
* ``reverse_id``, for excluding not only the edges in the current minibatch but
also their reverse edges according to the ID mapping in the argument
:attr:`reverse_eids`.
* ``reverse_types``, for excluding not only the edges in the current minibatch
but also their reverse edges stored in another type according to
the argument :attr:`reverse_etypes`.
* User-defined exclusion rule. It is a callable with edges in the current
minibatch as a single argument and should return the edges to be excluded.
degree_threshold : int
The threshold of node degrees, if the source or target node of an edge incident to
has larger or equal degrees than ``degree_threshold``, this edge will be excluded from
the graph
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.
reverse_etypes : dict[etype, etype], optional
The mapping from the original edge types to their reverse edge types.
Examples
--------
.. code:: python
low_degree_excluder = SpotTarget(g, degree_threshold=10)
sampler = as_edge_prediction_sampler(sampler, exclude=low_degree_excluder,
reverse_eids=reverse_eids, negative_sampler=negative_sampler.Uniform(1))
"""
def __init__(
self,
g,
exclude,
degree_threshold=10,
reverse_eids=None,
reverse_etypes=None,
):
self.g = g
self.exclude = exclude
self.degree_threshold = degree_threshold
self.reverse_eids = reverse_eids
self.reverse_etypes = reverse_etypes
def __call__(self, seed_edges):
g = self.g
src, dst = g.find_edges(seed_edges)
head_degree = g.in_degrees(src)
tail_degree = g.in_degrees(dst)
degree = torch.min(head_degree, tail_degree)
degree_mask = degree < self.degree_threshold
edges_need_to_exclude = seed_edges[degree_mask]
return find_exclude_eids(
g,
edges_need_to_exclude,
self.exclude,
self.reverse_eids,
self.reverse_etypes,
)