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/sampling/__init__.py
+14
View File
@@ -0,0 +1,14 @@
"""The ``dgl.sampling`` package contains operators and utilities for
sampling from a graph via random walks, neighbor sampling, etc. They
are typically used together with the ``DataLoader`` s in the
``dgl.dataloading`` package. The user guide :ref:`guide-minibatch`
gives a holistic explanation on how different components work together.
"""
from .randomwalks import *
from .pinsage import *
from .neighbor import *
from .labor import *
from .node2vec_randomwalk import *
from .negative import *
from . import utils
python/dgl/sampling/labor.py
+361
View File
@@ -0,0 +1,361 @@
#
# 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.py
#
"""Labor sampling APIs"""
from .. import backend as F, ndarray as nd, utils
from .._ffi.function import _init_api
from ..base import DGLError
from ..heterograph import DGLGraph
from ..random import choice
from .utils import EidExcluder
__all__ = ["sample_labors"]
def sample_labors(
g,
nodes,
fanout,
edge_dir="in",
prob=None,
importance_sampling=0,
random_seed=None,
seed2_contribution=0,
copy_ndata=True,
copy_edata=True,
exclude_edges=None,
output_device=None,
):
"""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.
For each node, a number of inbound (or outbound when ``edge_dir == 'out'``) edges
will be randomly chosen. The graph returned will then contain all the nodes in the
original graph, but only the sampled edges.
Node/edge features are not preserved. The original IDs of
the sampled edges are stored as the `dgl.EID` feature in the returned graph.
Parameters
----------
g : DGLGraph
The graph, allowed to have multiple node or edge types. Can be either on CPU or GPU.
nodes : tensor or dict
Node IDs to sample neighbors from.
This argument can take a single ID tensor or a dictionary of node types and ID tensors.
If a single tensor is given, the graph must only have one type of nodes.
fanout : int or dict[etype, int]
The number of edges to be sampled for each node on each edge type.
This argument can take a single int or a dictionary of edge types and ints.
If a single int is given, DGL will sample this number of edges for each node for
every edge type.
If -1 is given for a single edge type, all the neighboring edges with that edge
type will be selected.
edge_dir : str, optional
Determines whether to sample inbound or outbound edges.
Can take either ``in`` for inbound edges or ``out`` for outbound edges.
prob : str, optional
Feature name used as the (unnormalized) probabilities associated with each
neighboring edge of a node. The feature must have only one element for each
edge.
The features must be non-negative floats, and the sum of the features of
inbound/outbound edges for every node must be positive (though they don't have
to sum up to one). Otherwise, the result will be undefined.
If :attr:`prob` is not None, GPU sampling is not supported.
importance_sampling : int, optional
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.
random_seed : tensor
An int64 tensor with one element.
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 call to this function 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 a ``random_seed``, we get the random seed by getting a
random number from DGL. Use this argument with identical random_seed if multiple calls to
this function are used to sample as part of a single batch.
seed2_contribution : float, optional
A float value between [0, 1) that determines the contribution
of the second random seed to generate the random variates for the
LABOR sampling algorithm.
copy_ndata: bool, optional
If True, the node features of the new graph are copied from
the original graph. If False, the new graph will not have any
node features.
(Default: True)
copy_edata: bool, optional
If True, the edge features of the new graph are copied from
the original graph. If False, the new graph will not have any
edge features.
(Default: True)
exclude_edges: 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.
output_device : Framework-specific device context object, optional
The output device. Default is the same as the input graph.
Returns
-------
tuple(DGLGraph, list[Tensor])
A sampled subgraph containing only the sampled neighboring edges along with edge weights.
Notes
-----
If :attr:`copy_ndata` or :attr:`copy_edata` is True, same tensors are used as
the node or edge features of the original graph and the new graph.
As a result, users should avoid performing in-place operations
on the node features of the new graph to avoid feature corruption.
Examples
--------
Assume that you have the following graph
>>> g = dgl.graph(([0, 0, 1, 1, 2, 2], [1, 2, 0, 1, 2, 0]))
And the weights
>>> g.edata['prob'] = torch.FloatTensor([0., 1., 0., 1., 0., 1.])
To sample one inbound edge for node 0 and node 1:
>>> sg = dgl.sampling.sample_labors(g, [0, 1], 1)
>>> sg.edges(order='eid')
(tensor([1, 0]), tensor([0, 1]))
>>> sg.edata[dgl.EID]
tensor([2, 0])
To sample one inbound edge for node 0 and node 1 with probability in edge feature
``prob``:
>>> sg = dgl.sampling.sample_labors(g, [0, 1], 1, prob='prob')
>>> sg.edges(order='eid')
(tensor([2, 1]), tensor([0, 1]))
With ``fanout`` greater than the number of actual neighbors and without replacement,
DGL will take all neighbors instead:
>>> sg = dgl.sampling.sample_labors(g, [0, 1], 3)
>>> sg.edges(order='eid')
(tensor([1, 2, 0, 1]), tensor([0, 0, 1, 1]))
To exclude certain EID's during sampling for the seed nodes:
>>> g = dgl.graph(([0, 0, 1, 1, 2, 2], [1, 2, 0, 1, 2, 0]))
>>> g_edges = g.all_edges(form='all')``
(tensor([0, 0, 1, 1, 2, 2]), tensor([1, 2, 0, 1, 2, 0]), tensor([0, 1, 2, 3, 4, 5]))
>>> sg = dgl.sampling.sample_labors(g, [0, 1], 3, exclude_edges=[0, 1, 2])
>>> sg.all_edges(form='all')
(tensor([2, 1]), tensor([0, 1]), tensor([0, 1]))
>>> sg.has_edges_between(g_edges[0][:3],g_edges[1][:3])
tensor([False, False, False])
>>> g = dgl.heterograph({
... ('drug', 'interacts', 'drug'): ([0, 0, 1, 1, 3, 2], [1, 2, 0, 1, 2, 0]),
... ('drug', 'interacts', 'gene'): ([0, 0, 1, 1, 2, 2], [1, 2, 0, 1, 2, 0]),
... ('drug', 'treats', 'disease'): ([0, 0, 1, 1, 2, 2], [1, 2, 0, 1, 2, 0])})
>>> g_edges = g.all_edges(form='all', etype=('drug', 'interacts', 'drug'))
(tensor([0, 0, 1, 1, 3, 2]), tensor([1, 2, 0, 1, 2, 0]), tensor([0, 1, 2, 3, 4, 5]))
>>> excluded_edges = {('drug', 'interacts', 'drug'): g_edges[2][:3]}
>>> sg = dgl.sampling.sample_labors(g, {'drug':[0, 1]}, 3, exclude_edges=excluded_edges)
>>> sg.all_edges(form='all', etype=('drug', 'interacts', 'drug'))
(tensor([2, 1]), tensor([0, 1]), tensor([0, 1]))
>>> sg.has_edges_between(g_edges[0][:3],g_edges[1][:3],etype=('drug', 'interacts', 'drug'))
tensor([False, False, False])
"""
if F.device_type(g.device) == "cpu" and not g.is_pinned():
frontier, importances = _sample_labors(
g,
nodes,
fanout,
edge_dir=edge_dir,
prob=prob,
importance_sampling=importance_sampling,
random_seed=random_seed,
seed2_contribution=seed2_contribution,
copy_ndata=copy_ndata,
copy_edata=copy_edata,
exclude_edges=exclude_edges,
)
else:
frontier, importances = _sample_labors(
g,
nodes,
fanout,
edge_dir=edge_dir,
prob=prob,
importance_sampling=importance_sampling,
random_seed=random_seed,
seed2_contribution=seed2_contribution,
copy_ndata=copy_ndata,
copy_edata=copy_edata,
)
if exclude_edges is not None:
eid_excluder = EidExcluder(exclude_edges)
frontier, importances = eid_excluder(frontier, importances)
if output_device is None:
return (frontier, importances)
else:
return (
frontier.to(output_device),
list(map(lambda x: x.to(output_device), importances)),
)
def _sample_labors(
g,
nodes,
fanout,
edge_dir="in",
prob=None,
importance_sampling=0,
random_seed=None,
seed2_contribution=0,
copy_ndata=True,
copy_edata=True,
exclude_edges=None,
):
if random_seed is None:
random_seed = F.to_dgl_nd(choice(1e18, 1))
if not isinstance(nodes, dict):
if len(g.ntypes) > 1:
raise DGLError(
"Must specify node type when the graph is not homogeneous."
)
nodes = {g.ntypes[0]: nodes}
nodes = utils.prepare_tensor_dict(g, nodes, "nodes")
if len(nodes) == 0:
raise ValueError(
"Got an empty dictionary in the nodes argument. "
"Please pass in a dictionary with empty tensors as values instead."
)
ctx = utils.to_dgl_context(F.context(next(iter(nodes.values()))))
nodes_all_types = []
# nids_all_types is needed if one wants labor to work for subgraphs whose vertices have
# been renamed and the rolled randoms should be rolled for global vertex ids.
# It is disabled for now below by passing empty ndarrays.
nids_all_types = [nd.array([], ctx=ctx) for _ in g.ntypes]
for ntype in g.ntypes:
if ntype in nodes:
nodes_all_types.append(F.to_dgl_nd(nodes[ntype]))
else:
nodes_all_types.append(nd.array([], ctx=ctx))
if isinstance(fanout, nd.NDArray):
fanout_array = fanout
else:
if not isinstance(fanout, dict):
fanout_array = [int(fanout)] * len(g.etypes)
else:
if len(fanout) != len(g.etypes):
raise DGLError(
"Fan-out must be specified for each edge type "
"if a dict is provided."
)
fanout_array = [None] * len(g.etypes)
for etype, value in fanout.items():
fanout_array[g.get_etype_id(etype)] = value
fanout_array = F.to_dgl_nd(F.tensor(fanout_array, dtype=F.int64))
if (
isinstance(prob, list)
and len(prob) > 0
and isinstance(prob[0], nd.NDArray)
):
prob_arrays = prob
elif prob is None:
prob_arrays = [nd.array([], ctx=nd.cpu())] * len(g.etypes)
else:
prob_arrays = []
for etype in g.canonical_etypes:
if prob in g.edges[etype].data:
prob_arrays.append(F.to_dgl_nd(g.edges[etype].data[prob]))
else:
prob_arrays.append(nd.array([], ctx=nd.cpu()))
excluded_edges_all_t = []
if exclude_edges is not None:
if not isinstance(exclude_edges, dict):
if len(g.etypes) > 1:
raise DGLError(
"Must specify etype when the graph is not homogeneous."
)
exclude_edges = {g.canonical_etypes[0]: exclude_edges}
exclude_edges = utils.prepare_tensor_dict(g, exclude_edges, "edges")
for etype in g.canonical_etypes:
if etype in exclude_edges:
excluded_edges_all_t.append(F.to_dgl_nd(exclude_edges[etype]))
else:
excluded_edges_all_t.append(nd.array([], ctx=ctx))
ret_val = _CAPI_DGLSampleLabors(
g._graph,
nodes_all_types,
fanout_array,
edge_dir,
prob_arrays,
excluded_edges_all_t,
importance_sampling,
random_seed,
seed2_contribution,
nids_all_types,
)
subgidx = ret_val[0]
importances = [F.from_dgl_nd(importance) for importance in ret_val[1:]]
induced_edges = subgidx.induced_edges
ret = DGLGraph(subgidx.graph, g.ntypes, g.etypes)
if copy_ndata:
node_frames = utils.extract_node_subframes(g, None)
utils.set_new_frames(ret, node_frames=node_frames)
if copy_edata:
edge_frames = utils.extract_edge_subframes(g, induced_edges)
utils.set_new_frames(ret, edge_frames=edge_frames)
return ret, importances
DGLGraph.sample_labors = utils.alias_func(sample_labors)
_init_api("dgl.sampling.labor", __name__)
python/dgl/sampling/negative.py
+126
View File
@@ -0,0 +1,126 @@
"""Negative sampling APIs"""
from numpy.polynomial import polynomial
from .. import backend as F, utils
from .._ffi.function import _init_api
from ..heterograph import DGLGraph
__all__ = ["global_uniform_negative_sampling"]
def _calc_redundancy(
k_hat, num_edges, num_pairs, r=3
): # pylint: disable=invalid-name
# pylint: disable=invalid-name
# Calculates the number of samples required based on a lower-bound
# of the expected number of negative samples, based on N draws from
# a binomial distribution. Solves the following equation for N:
#
# k_hat = N*p_k - r * np.sqrt(N*p_k*(1-p_k))
#
# where p_k is the probability that a node pairing is a negative edge
# and r is the number of standard deviations to construct the lower bound
#
# Credits to @zjost
p_m = num_edges / num_pairs
p_k = 1 - p_m
a = p_k**2
b = -p_k * (2 * k_hat + r**2 * p_m)
c = k_hat**2
poly = polynomial.Polynomial([c, b, a])
N = poly.roots()[-1]
redundancy = N / k_hat - 1.0
return redundancy
def global_uniform_negative_sampling(
g,
num_samples,
exclude_self_loops=True,
replace=False,
etype=None,
redundancy=None,
):
"""Performs negative sampling, which generate source-destination pairs such that
edges with the given type do not exist.
Specifically, this function takes in an edge type and a number of samples. It
returns two tensors ``src`` and ``dst``, the former in the range of ``[0, num_src)``
and the latter in the range of ``[0, num_dst)``, where ``num_src`` and ``num_dst``
represents the number of nodes with the source and destination node type respectively.
It guarantees that no edge will exist between the corresponding pairs of ``src``
with the source node type and ``dst`` with the destination node type.
.. note::
This negative sampler will try to generate as many negative samples as possible, but
it may rarely return less than :attr:`num_samples` negative samples.
This is more likely to happen when a graph is so small or dense that not many
unique negative samples exist.
Parameters
----------
g : DGLGraph
The graph.
num_samples : int
The number of desired negative samples to generate.
exclude_self_loops : bool, optional
Whether to exclude self-loops from the negative samples. Only impacts the
edge types whose source and destination node types are the same.
Default: True.
replace : bool, optional
Whether to sample with replacement. Setting it to True will make things
faster. (Default: False)
etype : str or tuple of str, optional
The edge type. Can be omitted if the graph only has one edge type.
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:`num_samples` negative
samples, but will also take more time and memory.
(Default: automatically determined by the density of graph)
Returns
-------
tuple[Tensor, Tensor]
The source and destination pairs.
Examples
--------
>>> g = dgl.graph(([0, 1, 2], [1, 2, 3]))
>>> dgl.sampling.global_uniform_negative_sampling(g, 3)
(tensor([0, 1, 3]), tensor([2, 0, 2]))
"""
if etype is None:
etype = g.etypes[0]
utype, _, vtype = g.to_canonical_etype(etype)
exclude_self_loops = exclude_self_loops and (utype == vtype)
redundancy = _calc_redundancy(
num_samples, g.num_edges(etype), g.num_nodes(utype) * g.num_nodes(vtype)
)
etype_id = g.get_etype_id(etype)
src, dst = _CAPI_DGLGlobalUniformNegativeSampling(
g._graph,
etype_id,
num_samples,
3,
exclude_self_loops,
replace,
redundancy,
)
return F.from_dgl_nd(src), F.from_dgl_nd(dst)
DGLGraph.global_uniform_negative_sampling = utils.alias_func(
global_uniform_negative_sampling
)
_init_api("dgl.sampling.negative", __name__)
python/dgl/sampling/neighbor.py
+1039
View File
File diff suppressed because it is too large Load Diff
python/dgl/sampling/node2vec_randomwalk.py
+103
View File
@@ -0,0 +1,103 @@
"""Node2vec random walk"""
from .. import backend as F, ndarray as nd, utils
from .._ffi.function import _init_api
# pylint: disable=invalid-name
__all__ = ["node2vec_random_walk"]
def node2vec_random_walk(
g, nodes, p, q, walk_length, prob=None, return_eids=False
):
"""
Generate random walk traces from an array of starting nodes based on the node2vec model.
Paper: `node2vec: Scalable Feature Learning for Networks
<https://arxiv.org/abs/1607.00653>`__.
The returned traces all have length ``walk_length + 1``, where the first node
is the starting node itself.
Note that if a random walk stops in advance, DGL pads the trace with -1 to have the same
length.
Parameters
----------
g : DGLGraph
The graph. Must be on CPU.
Note that node2vec only support homogeneous graph.
nodes : Tensor
Node ID tensor from which the random walk traces starts.
The tensor must be on CPU, and must have the same dtype as the ID type
of the graph.
p: float
Likelihood of immediately revisiting a node in the walk.
q: float
Control parameter to interpolate between breadth-first strategy and depth-first strategy.
walk_length: int
Length of random walks.
prob : str, optional
The name of the edge feature tensor on the graph storing the (unnormalized)
probabilities associated with each edge for choosing the next node.
The feature tensor must be non-negative and the sum of the probabilities
must be positive for the outbound edges of all nodes (although they don't have
to sum up to one). The result will be undefined otherwise.
If omitted, DGL assumes that the neighbors are picked uniformly.
return_eids : bool, optional
If True, additionally return the edge IDs traversed.
Default: False.
Returns
-------
traces : Tensor
A 2-dimensional node ID tensor with shape ``(num_seeds, walk_length + 1)``.
eids : Tensor, optional
A 2-dimensional edge ID tensor with shape ``(num_seeds, length)``.
Only returned if :attr:`return_eids` is True.
Examples
--------
>>> g1 = dgl.graph(([0, 1, 1, 2, 3], [1, 2, 3, 0, 0]))
>>> dgl.sampling.node2vec_random_walk(g1, [0, 1, 2, 0], 1, 1, walk_length=4)
tensor([[0, 1, 3, 0, 1],
[1, 2, 0, 1, 3],
[2, 0, 1, 3, 0],
[0, 1, 2, 0, 1]])
>>> dgl.sampling.node2vec_random_walk(g1, [0, 1, 2, 0], 1, 1, walk_length=4, return_eids=True)
(tensor([[0, 1, 3, 0, 1],
[1, 2, 0, 1, 2],
[2, 0, 1, 2, 0],
[0, 1, 2, 0, 1]]),
tensor([[0, 2, 4, 0],
[1, 3, 0, 1],
[3, 0, 1, 3],
[0, 1, 3, 0]]))
"""
assert g.device == F.cpu(), "Graph must be on CPU."
gidx = g._graph
nodes = F.to_dgl_nd(utils.prepare_tensor(g, nodes, "nodes"))
if prob is None:
prob_nd = nd.array([], ctx=nodes.ctx)
else:
prob_nd = F.to_dgl_nd(g.edata[prob])
traces, eids = _CAPI_DGLSamplingNode2vec(
gidx, nodes, p, q, walk_length, prob_nd
)
traces = F.from_dgl_nd(traces)
eids = F.from_dgl_nd(eids)
return (traces, eids) if return_eids else traces
_init_api("dgl.sampling.randomwalks", __name__)
python/dgl/sampling/pinsage.py
+275
View File
@@ -0,0 +1,275 @@
"""PinSAGE sampler & related functions and classes"""
import numpy as np
from .. import backend as F, convert, utils
from .._ffi.function import _init_api
from .randomwalks import random_walk
def _select_pinsage_neighbors(src, dst, num_samples_per_node, k):
"""Determine the neighbors for PinSAGE algorithm from the given random walk traces.
This is fusing ``to_simple()``, ``select_topk()``, and counting the number of occurrences
together.
"""
src = F.to_dgl_nd(src)
dst = F.to_dgl_nd(dst)
src, dst, counts = _CAPI_DGLSamplingSelectPinSageNeighbors(
src, dst, num_samples_per_node, k
)
src = F.from_dgl_nd(src)
dst = F.from_dgl_nd(dst)
counts = F.from_dgl_nd(counts)
return (src, dst, counts)
class RandomWalkNeighborSampler(object):
"""PinSage-like neighbor sampler extended to any heterogeneous graphs.
Given a heterogeneous graph and a list of nodes, this callable will generate a homogeneous
graph where the neighbors of each given node are the most commonly visited nodes of the
same type by multiple random walks starting from that given node. Each random walk consists
of multiple metapath-based traversals, with a probability of termination after each traversal.
The edges of the returned homogeneous graph will connect to the given nodes from their most
commonly visited nodes, with a feature indicating the number of visits.
The metapath must have the same beginning and ending node type to make the algorithm work.
This is a generalization of PinSAGE sampler which only works on bidirectional bipartite
graphs.
UVA and GPU sampling is supported for this sampler.
Refer to :ref:`guide-minibatch-gpu-sampling` for more details.
Parameters
----------
G : DGLGraph
The graph.
num_traversals : int
The maximum number of metapath-based traversals for a single random walk.
Usually considered a hyperparameter.
termination_prob : float
Termination probability after each metapath-based traversal.
Usually considered a hyperparameter.
num_random_walks : int
Number of random walks to try for each given node.
Usually considered a hyperparameter.
num_neighbors : int
Number of neighbors (or most commonly visited nodes) to select for each given node.
metapath : list[str] or list[tuple[str, str, str]], optional
The metapath.
If not given, DGL assumes that the graph is homogeneous and the metapath consists
of one step over the single edge type.
weight_column : str, default "weights"
The name of the edge feature to be stored on the returned graph with the number of
visits.
Examples
--------
See examples in :any:`PinSAGESampler`.
"""
def __init__(
self,
G,
num_traversals,
termination_prob,
num_random_walks,
num_neighbors,
metapath=None,
weight_column="weights",
):
self.G = G
self.weight_column = weight_column
self.num_random_walks = num_random_walks
self.num_neighbors = num_neighbors
self.num_traversals = num_traversals
if metapath is None:
if len(G.ntypes) > 1 or len(G.etypes) > 1:
raise ValueError(
"Metapath must be specified if the graph is homogeneous."
)
metapath = [G.canonical_etypes[0]]
start_ntype = G.to_canonical_etype(metapath[0])[0]
end_ntype = G.to_canonical_etype(metapath[-1])[-1]
if start_ntype != end_ntype:
raise ValueError(
"The metapath must start and end at the same node type."
)
self.ntype = start_ntype
self.metapath_hops = len(metapath)
self.metapath = metapath
self.full_metapath = metapath * num_traversals
restart_prob = np.zeros(self.metapath_hops * num_traversals)
restart_prob[
self.metapath_hops :: self.metapath_hops
] = termination_prob
restart_prob = F.tensor(restart_prob, dtype=F.float32)
self.restart_prob = F.copy_to(restart_prob, G.device)
# pylint: disable=no-member
def __call__(self, seed_nodes):
"""
Parameters
----------
seed_nodes : Tensor
A tensor of given node IDs of node type ``ntype`` to generate neighbors from. The
node type ``ntype`` is the beginning and ending node type of the given metapath.
It must be on the same device as the graph and have the same dtype
as the ID type of the graph.
Returns
-------
g : DGLGraph
A homogeneous graph constructed by selecting neighbors for each given node according
to the algorithm above.
"""
seed_nodes = utils.prepare_tensor(self.G, seed_nodes, "seed_nodes")
self.restart_prob = F.copy_to(self.restart_prob, F.context(seed_nodes))
seed_nodes = F.repeat(seed_nodes, self.num_random_walks, 0)
paths, _ = random_walk(
self.G,
seed_nodes,
metapath=self.full_metapath,
restart_prob=self.restart_prob,
)
src = F.reshape(
paths[:, self.metapath_hops :: self.metapath_hops], (-1,)
)
dst = F.repeat(paths[:, 0], self.num_traversals, 0)
src, dst, counts = _select_pinsage_neighbors(
src,
dst,
(self.num_random_walks * self.num_traversals),
self.num_neighbors,
)
neighbor_graph = convert.heterograph(
{(self.ntype, "_E", self.ntype): (src, dst)},
{self.ntype: self.G.num_nodes(self.ntype)},
)
neighbor_graph.edata[self.weight_column] = counts
return neighbor_graph
class PinSAGESampler(RandomWalkNeighborSampler):
"""PinSAGE-like neighbor sampler.
This callable works on a bidirectional bipartite graph with edge types
``(ntype, fwtype, other_type)`` and ``(other_type, bwtype, ntype)`` (where ``ntype``,
``fwtype``, ``bwtype`` and ``other_type`` could be arbitrary type names). It will generate
a homogeneous graph of node type ``ntype`` where the neighbors of each given node are the
most commonly visited nodes of the same type by multiple random walks starting from that
given node. Each random walk consists of multiple metapath-based traversals, with a
probability of termination after each traversal. The metapath is always ``[fwtype, bwtype]``,
walking from node type ``ntype`` to node type ``other_type`` then back to ``ntype``.
The edges of the returned homogeneous graph will connect to the given nodes from their most
commonly visited nodes, with a feature indicating the number of visits.
UVA and GPU sampling is supported for this sampler.
Refer to :ref:`guide-minibatch-gpu-sampling` for more details.
Parameters
----------
G : DGLGraph
The bidirectional bipartite graph.
The graph should only have two node types: ``ntype`` and ``other_type``.
The graph should only have two edge types, one connecting from ``ntype`` to
``other_type``, and another connecting from ``other_type`` to ``ntype``.
ntype : str
The node type for which the graph would be constructed on.
other_type : str
The other node type.
num_traversals : int
The maximum number of metapath-based traversals for a single random walk.
Usually considered a hyperparameter.
termination_prob : int
Termination probability after each metapath-based traversal.
Usually considered a hyperparameter.
num_random_walks : int
Number of random walks to try for each given node.
Usually considered a hyperparameter.
num_neighbors : int
Number of neighbors (or most commonly visited nodes) to select for each given node.
weight_column : str, default "weights"
The name of the edge feature to be stored on the returned graph with the number of
visits.
Examples
--------
Generate a random bidirectional bipartite graph with 3000 "A" nodes and 5000 "B" nodes.
>>> g = scipy.sparse.random(3000, 5000, 0.003)
>>> G = dgl.heterograph({
... ('A', 'AB', 'B'): g.nonzero(),
... ('B', 'BA', 'A'): g.T.nonzero()})
Then we create a PinSage neighbor sampler that samples a graph of node type "A". Each
node would have (a maximum of) 10 neighbors.
>>> sampler = dgl.sampling.PinSAGESampler(G, 'A', 'B', 3, 0.5, 200, 10)
This is how we select the neighbors for node #0, #1 and #2 of type "A" according to
PinSAGE algorithm:
>>> seeds = torch.LongTensor([0, 1, 2])
>>> frontier = sampler(seeds)
>>> frontier.all_edges(form='uv')
(tensor([ 230, 0, 802, 47, 50, 1639, 1533, 406, 2110, 2687, 2408, 2823,
0, 972, 1230, 1658, 2373, 1289, 1745, 2918, 1818, 1951, 1191, 1089,
1282, 566, 2541, 1505, 1022, 812]),
tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2]))
For an end-to-end example of PinSAGE model, including sampling on multiple layers
and computing with the sampled graphs, please refer to our PinSage example
in ``examples/pytorch/pinsage``.
References
----------
Graph Convolutional Neural Networks for Web-Scale Recommender Systems
Ying et al., 2018, https://arxiv.org/abs/1806.01973
"""
def __init__(
self,
G,
ntype,
other_type,
num_traversals,
termination_prob,
num_random_walks,
num_neighbors,
weight_column="weights",
):
metagraph = G.metagraph()
fw_etype = list(metagraph[ntype][other_type])[0]
bw_etype = list(metagraph[other_type][ntype])[0]
super().__init__(
G,
num_traversals,
termination_prob,
num_random_walks,
num_neighbors,
metapath=[fw_etype, bw_etype],
weight_column=weight_column,
)
_init_api("dgl.sampling.pinsage", __name__)
python/dgl/sampling/randomwalks.py
+314
View File
@@ -0,0 +1,314 @@
"""Random walk routines
"""
from .. import backend as F, ndarray as nd, utils
from .._ffi.function import _init_api
from ..base import DGLError
__all__ = ["random_walk", "pack_traces"]
def random_walk(
g,
nodes,
*,
metapath=None,
length=None,
prob=None,
restart_prob=None,
return_eids=False
):
"""Generate random walk traces from an array of starting nodes based on the given metapath.
Each starting node will have one trace generated, which
1. Start from the given node and set ``t`` to 0.
2. Pick and traverse along edge type ``metapath[t]`` from the current node.
3. If no edge can be found, halt. Otherwise, increment ``t`` and go to step 2.
To generate multiple traces for a single node, you can specify the same node multiple
times.
The returned traces all have length ``len(metapath) + 1``, where the first node
is the starting node itself.
If a random walk stops in advance, DGL pads the trace with -1 to have the same
length.
This function supports the graph on GPU and UVA sampling.
Parameters
----------
g : DGLGraph
The graph.
nodes : Tensor
Node ID tensor from which the random walk traces starts.
The tensor must have the same dtype as the ID type of the graph.
The tensor must be on the same device as the graph or
on the GPU when the graph is pinned (UVA sampling).
metapath : list[str or tuple of str], optional
Metapath, specified as a list of edge types.
Mutually exclusive with :attr:`length`.
If omitted, DGL assumes that ``g`` only has one node & edge type. In this
case, the argument ``length`` specifies the length of random walk traces.
length : int, optional
Length of random walks.
Mutually exclusive with :attr:`metapath`.
Only used when :attr:`metapath` is None.
prob : str, optional
The name of the edge feature tensor on the graph storing the (unnormalized)
probabilities associated with each edge for choosing the next node.
The feature tensor must be non-negative and the sum of the probabilities
must be positive for the outbound edges of all nodes (although they don't have
to sum up to one). The result will be undefined otherwise.
The feature tensor must be on the same device as the graph.
If omitted, DGL assumes that the neighbors are picked uniformly.
restart_prob : float or Tensor, optional
Probability to terminate the current trace before each transition.
If a tensor is given, :attr:`restart_prob` should be on the same device as the graph
or on the GPU when the graph is pinned (UVA sampling),
and have the same length as :attr:`metapath` or :attr:`length`.
return_eids : bool, optional
If True, additionally return the edge IDs traversed.
Default: False.
Returns
-------
traces : Tensor
A 2-dimensional node ID tensor with shape ``(num_seeds, len(metapath) + 1)`` or
``(num_seeds, length + 1)`` if :attr:`metapath` is None.
eids : Tensor, optional
A 2-dimensional edge ID tensor with shape ``(num_seeds, len(metapath))`` or
``(num_seeds, length)`` if :attr:`metapath` is None. Only returned if
:attr:`return_eids` is True.
types : Tensor
A 1-dimensional node type ID tensor with shape ``(len(metapath) + 1)`` or
``(length + 1)``.
The type IDs match the ones in the original graph ``g``.
Examples
--------
The following creates a homogeneous graph:
>>> g1 = dgl.graph(([0, 1, 1, 2, 3], [1, 2, 3, 0, 0]))
Normal random walk:
>>> dgl.sampling.random_walk(g1, [0, 1, 2, 0], length=4)
(tensor([[0, 1, 2, 0, 1],
[1, 3, 0, 1, 3],
[2, 0, 1, 3, 0],
[0, 1, 2, 0, 1]]), tensor([0, 0, 0, 0, 0]))
Or returning edge IDs:
>>> dgl.sampling.random_walk(g1, [0, 1, 2, 0], length=4, return_eids=True)
(tensor([[0, 1, 2, 0, 1],
[1, 3, 0, 1, 2],
[2, 0, 1, 3, 0],
[0, 1, 3, 0, 1]]),
tensor([[0, 1, 3, 0],
[2, 4, 0, 1],
[3, 0, 2, 4],
[0, 2, 4, 0]]),
tensor([0, 0, 0, 0, 0]))
The first tensor indicates the random walk path for each seed node.
The j-th element in the second tensor indicates the node type ID of the j-th node
in every path. In this case, it is returning all 0.
Random walk with restart:
>>> dgl.sampling.random_walk_with_restart(g1, [0, 1, 2, 0], length=4, restart_prob=0.5)
(tensor([[ 0, -1, -1, -1, -1],
[ 1, 3, 0, -1, -1],
[ 2, -1, -1, -1, -1],
[ 0, -1, -1, -1, -1]]), tensor([0, 0, 0, 0, 0]))
Non-uniform random walk:
>>> g1.edata['p'] = torch.FloatTensor([1, 0, 1, 1, 1]) # disallow going from 1 to 2
>>> dgl.sampling.random_walk(g1, [0, 1, 2, 0], length=4, prob='p')
(tensor([[0, 1, 3, 0, 1],
[1, 3, 0, 1, 3],
[2, 0, 1, 3, 0],
[0, 1, 3, 0, 1]]), tensor([0, 0, 0, 0, 0]))
Metapath-based random walk:
>>> g2 = dgl.heterograph({
... ('user', 'follow', 'user'): ([0, 1, 1, 2, 3], [1, 2, 3, 0, 0]),
... ('user', 'view', 'item'): ([0, 0, 1, 2, 3, 3], [0, 1, 1, 2, 2, 1]),
... ('item', 'viewed-by', 'user'): ([0, 1, 1, 2, 2, 1], [0, 0, 1, 2, 3, 3])
>>> dgl.sampling.random_walk(
... g2, [0, 1, 2, 0], metapath=['follow', 'view', 'viewed-by'] * 2)
(tensor([[0, 1, 1, 1, 2, 2, 3],
[1, 3, 1, 1, 2, 2, 2],
[2, 0, 1, 1, 3, 1, 1],
[0, 1, 1, 0, 1, 1, 3]]), tensor([0, 0, 1, 0, 0, 1, 0]))
Metapath-based random walk, with restarts only on items (i.e. after traversing a "view"
relationship):
>>> dgl.sampling.random_walk(
... g2, [0, 1, 2, 0], metapath=['follow', 'view', 'viewed-by'] * 2,
... restart_prob=torch.FloatTensor([0, 0.5, 0, 0, 0.5, 0]))
(tensor([[ 0, 1, -1, -1, -1, -1, -1],
[ 1, 3, 1, 0, 1, 1, 0],
[ 2, 0, 1, 1, 3, 2, 2],
[ 0, 1, 1, 3, 0, 0, 0]]), tensor([0, 0, 1, 0, 0, 1, 0]))
"""
n_etypes = len(g.canonical_etypes)
n_ntypes = len(g.ntypes)
if metapath is None:
if n_etypes > 1 or n_ntypes > 1:
raise DGLError(
"metapath not specified and the graph is not homogeneous."
)
if length is None:
raise ValueError(
"Please specify either the metapath or the random walk length."
)
metapath = [0] * length
else:
metapath = [g.get_etype_id(etype) for etype in metapath]
gidx = g._graph
nodes = utils.prepare_tensor(g, nodes, "nodes")
nodes = F.to_dgl_nd(nodes)
# (Xin) Since metapath array is created by us, safe to skip the check
# and keep it on CPU to make max_nodes sanity check easier.
metapath = F.to_dgl_nd(F.astype(F.tensor(metapath), g.idtype))
# Load the probability tensor from the edge frames
ctx = utils.to_dgl_context(g.device)
if prob is None:
p_nd = [nd.array([], ctx=ctx) for _ in g.canonical_etypes]
else:
p_nd = []
for etype in g.canonical_etypes:
if prob in g.edges[etype].data:
prob_nd = F.to_dgl_nd(g.edges[etype].data[prob])
else:
prob_nd = nd.array([], ctx=ctx)
p_nd.append(prob_nd)
# Actual random walk
if restart_prob is None:
traces, eids, types = _CAPI_DGLSamplingRandomWalk(
gidx, nodes, metapath, p_nd
)
elif F.is_tensor(restart_prob):
restart_prob = F.to_dgl_nd(restart_prob)
traces, eids, types = _CAPI_DGLSamplingRandomWalkWithStepwiseRestart(
gidx, nodes, metapath, p_nd, restart_prob
)
elif isinstance(restart_prob, float):
traces, eids, types = _CAPI_DGLSamplingRandomWalkWithRestart(
gidx, nodes, metapath, p_nd, restart_prob
)
else:
raise TypeError("restart_prob should be float or Tensor.")
traces = F.from_dgl_nd(traces)
types = F.from_dgl_nd(types)
eids = F.from_dgl_nd(eids)
return (traces, eids, types) if return_eids else (traces, types)
def pack_traces(traces, types):
"""Pack the padded traces returned by ``random_walk()`` into a concatenated array.
The padding values (-1) are removed, and the length and offset of each trace is
returned along with the concatenated node ID and node type arrays.
Parameters
----------
traces : Tensor
A 2-dimensional node ID tensor. Must be on CPU and either ``int32`` or ``int64``.
types : Tensor
A 1-dimensional node type ID tensor. Must be on CPU and either ``int32`` or ``int64``.
Returns
-------
concat_vids : Tensor
An array of all node IDs concatenated and padding values removed.
concat_types : Tensor
An array of node types corresponding for each node in ``concat_vids``.
Has the same length as ``concat_vids``.
lengths : Tensor
Length of each trace in the original traces tensor.
offsets : Tensor
Offset of each trace in the originial traces tensor in the new concatenated tensor.
Notes
-----
The returned tensors are on CPU.
Examples
--------
>>> g2 = dgl.heterograph({
... ('user', 'follow', 'user'): ([0, 1, 1, 2, 3], [1, 2, 3, 0, 0]),
... ('user', 'view', 'item'): ([0, 0, 1, 2, 3, 3], [0, 1, 1, 2, 2, 1]),
... ('item', 'viewed-by', 'user'): ([0, 1, 1, 2, 2, 1], [0, 0, 1, 2, 3, 3])
>>> traces, types = dgl.sampling.random_walk(
... g2, [0, 0], metapath=['follow', 'view', 'viewed-by'] * 2,
... restart_prob=torch.FloatTensor([0, 0.5, 0, 0, 0.5, 0]))
>>> traces, types
(tensor([[ 0, 1, -1, -1, -1, -1, -1],
[ 0, 1, 1, 3, 0, 0, 0]]), tensor([0, 0, 1, 0, 0, 1, 0]))
>>> concat_vids, concat_types, lengths, offsets = dgl.sampling.pack_traces(traces, types)
>>> concat_vids
tensor([0, 1, 0, 1, 1, 3, 0, 0, 0])
>>> concat_types
tensor([0, 0, 0, 0, 1, 0, 0, 1, 0])
>>> lengths
tensor([2, 7])
>>> offsets
tensor([0, 2]))
The first tensor ``concat_vids`` is the concatenation of all paths, i.e. flattened array
of ``traces``, excluding all padding values (-1).
The second tensor ``concat_types`` stands for the node type IDs of all corresponding nodes
in the first tensor.
The third and fourth tensor indicates the length and the offset of each path. With these
tensors it is easy to obtain the i-th random walk path with:
>>> vids = concat_vids.split(lengths.tolist())
>>> vtypes = concat_vtypes.split(lengths.tolist())
>>> vids[1], vtypes[1]
(tensor([0, 1, 1, 3, 0, 0, 0]), tensor([0, 0, 1, 0, 0, 1, 0]))
"""
assert (
F.is_tensor(traces) and F.context(traces) == F.cpu()
), "traces must be a CPU tensor"
assert (
F.is_tensor(types) and F.context(types) == F.cpu()
), "types must be a CPU tensor"
traces = F.to_dgl_nd(traces)
types = F.to_dgl_nd(types)
concat_vids, concat_types, lengths, offsets = _CAPI_DGLSamplingPackTraces(
traces, types
)
concat_vids = F.from_dgl_nd(concat_vids)
concat_types = F.from_dgl_nd(concat_types)
lengths = F.from_dgl_nd(lengths)
offsets = F.from_dgl_nd(offsets)
return concat_vids, concat_types, lengths, offsets
_init_api("dgl.sampling.randomwalks", __name__)
python/dgl/sampling/utils.py
+105
View File
@@ -0,0 +1,105 @@
"""Sampling utilities"""
from collections.abc import Mapping
import numpy as np
from .. import backend as F, transforms, utils
from ..base import EID
from ..utils import recursive_apply, recursive_apply_pair
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 not isinstance(frontier_parent_eids, Mapping):
return np.isin(frontier_parent_eids, exclude_eids).nonzero()[0]
result = {}
for k, v in frontier_parent_eids.items():
if k in exclude_eids:
result[k] = np.isin(v, exclude_eids[k]).nonzero()[0]
return recursive_apply(result, F.zerocopy_from_numpy)
class EidExcluder(object):
"""Class that finds the edges whose IDs in parent graph appeared in exclude_eids.
The edge IDs can be both CPU and GPU tensors.
"""
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 = (
recursive_apply(exclude_eids, F.zerocopy_to_numpy)
if exclude_eids is not None
else None
)
else:
self._filter = recursive_apply(exclude_eids, utils.Filter)
def _find_indices(self, parent_eids):
"""Find the set of edge indices to remove."""
if self._exclude_eids is not None:
parent_eids_np = recursive_apply(parent_eids, F.zerocopy_to_numpy)
return _locate_eids_to_exclude(parent_eids_np, self._exclude_eids)
else:
assert self._filter is not None
func = lambda x, y: x.find_included_indices(y)
return recursive_apply_pair(self._filter, parent_eids, func)
def __call__(self, frontier, weights=None):
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 = transforms.remove_edges(
frontier, located_eids, store_ids=True
)
if (
weights is not None
and weights[0].shape[0] == frontier.num_edges()
):
weights[0] = F.gather_row(weights[0], frontier.edata[EID])
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 i, (k, v) in enumerate(located_eids.items()):
if len(v) > 0:
frontier = transforms.remove_edges(
frontier, v, etype=k, store_ids=True
)
new_eids[k] = F.gather_row(
parent_eids[k], frontier.edges[k].data[EID]
)
if weights is not None and weights[i].shape[
0
] == frontier.num_edges(k):
weights[i] = F.gather_row(
weights[i], frontier.edges[k].data[EID]
)
frontier.edata[EID] = new_eids
return frontier if weights is None else (frontier, weights)