dmlc--dgl
yxy235
389 行
16 KiB
Python
389 行
16 KiB
Python
"""Subgraph samplers"""
|
|
|
|
from collections import defaultdict
|
|
from typing import Dict
|
|
|
|
import torch
|
|
from torch.utils.data import functional_datapipe
|
|
|
|
from .base import etype_str_to_tuple
|
|
from .internal import compact_temporal_nodes, unique_and_compact
|
|
from .minibatch_transformer import MiniBatchTransformer
|
|
|
|
__all__ = [
|
|
"SubgraphSampler",
|
|
]
|
|
|
|
|
|
@functional_datapipe("sample_subgraph")
|
|
class SubgraphSampler(MiniBatchTransformer):
|
|
"""A subgraph sampler used to sample a subgraph from a given set of nodes
|
|
from a larger graph.
|
|
|
|
Functional name: :obj:`sample_subgraph`.
|
|
|
|
This class is the base class of all subgraph samplers. Any subclass of
|
|
SubgraphSampler should implement either the :meth:`sample_subgraphs` method
|
|
or the :meth:`sampling_stages` method to define the fine-grained sampling
|
|
stages to take advantage of optimizations provided by the GraphBolt
|
|
DataLoader.
|
|
|
|
Parameters
|
|
----------
|
|
datapipe : DataPipe
|
|
The datapipe.
|
|
args : Non-Keyword Arguments
|
|
Arguments to be passed into sampling_stages.
|
|
kwargs : Keyword Arguments
|
|
Arguments to be passed into sampling_stages.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
datapipe,
|
|
*args,
|
|
**kwargs,
|
|
):
|
|
datapipe = datapipe.transform(self._preprocess)
|
|
datapipe = self.sampling_stages(datapipe, *args, **kwargs)
|
|
datapipe = datapipe.transform(self._postprocess)
|
|
super().__init__(datapipe)
|
|
|
|
@staticmethod
|
|
def _postprocess(minibatch):
|
|
delattr(minibatch, "_seed_nodes")
|
|
delattr(minibatch, "_seeds_timestamp")
|
|
return minibatch
|
|
|
|
@staticmethod
|
|
def _preprocess(minibatch):
|
|
if minibatch.node_pairs is not None:
|
|
(
|
|
seeds,
|
|
seeds_timestamp,
|
|
minibatch.compacted_node_pairs,
|
|
minibatch.compacted_negative_srcs,
|
|
minibatch.compacted_negative_dsts,
|
|
) = SubgraphSampler._node_pairs_preprocess(minibatch)
|
|
elif minibatch.seed_nodes is not None:
|
|
seeds = minibatch.seed_nodes
|
|
seeds_timestamp = (
|
|
minibatch.timestamp if hasattr(minibatch, "timestamp") else None
|
|
)
|
|
elif minibatch.seeds is not None:
|
|
(
|
|
seeds,
|
|
seeds_timestamp,
|
|
minibatch.compacted_seeds,
|
|
) = SubgraphSampler._seeds_preprocess(minibatch)
|
|
else:
|
|
raise ValueError(
|
|
f"Invalid minibatch {minibatch}: One of `node_pairs`, "
|
|
"`seed_nodes` and `seeds` should have a value."
|
|
)
|
|
minibatch._seed_nodes = seeds
|
|
minibatch._seeds_timestamp = seeds_timestamp
|
|
return minibatch
|
|
|
|
@staticmethod
|
|
def _node_pairs_preprocess(minibatch):
|
|
use_timestamp = hasattr(minibatch, "timestamp")
|
|
node_pairs = minibatch.node_pairs
|
|
neg_src, neg_dst = minibatch.negative_srcs, minibatch.negative_dsts
|
|
has_neg_src = neg_src is not None
|
|
has_neg_dst = neg_dst is not None
|
|
is_heterogeneous = isinstance(node_pairs, Dict)
|
|
if is_heterogeneous:
|
|
has_neg_src = has_neg_src and all(
|
|
item is not None for item in neg_src.values()
|
|
)
|
|
has_neg_dst = has_neg_dst and all(
|
|
item is not None for item in neg_dst.values()
|
|
)
|
|
# Collect nodes from all types of input.
|
|
nodes = defaultdict(list)
|
|
nodes_timestamp = None
|
|
if use_timestamp:
|
|
nodes_timestamp = defaultdict(list)
|
|
for etype, (src, dst) in node_pairs.items():
|
|
src_type, _, dst_type = etype_str_to_tuple(etype)
|
|
nodes[src_type].append(src)
|
|
nodes[dst_type].append(dst)
|
|
if use_timestamp:
|
|
nodes_timestamp[src_type].append(minibatch.timestamp[etype])
|
|
nodes_timestamp[dst_type].append(minibatch.timestamp[etype])
|
|
if has_neg_src:
|
|
for etype, src in neg_src.items():
|
|
src_type, _, _ = etype_str_to_tuple(etype)
|
|
nodes[src_type].append(src.view(-1))
|
|
if use_timestamp:
|
|
nodes_timestamp[src_type].append(
|
|
minibatch.timestamp[etype].repeat_interleave(
|
|
src.shape[-1]
|
|
)
|
|
)
|
|
if has_neg_dst:
|
|
for etype, dst in neg_dst.items():
|
|
_, _, dst_type = etype_str_to_tuple(etype)
|
|
nodes[dst_type].append(dst.view(-1))
|
|
if use_timestamp:
|
|
nodes_timestamp[dst_type].append(
|
|
minibatch.timestamp[etype].repeat_interleave(
|
|
dst.shape[-1]
|
|
)
|
|
)
|
|
# Unique and compact the collected nodes.
|
|
if use_timestamp:
|
|
seeds, nodes_timestamp, compacted = compact_temporal_nodes(
|
|
nodes, nodes_timestamp
|
|
)
|
|
else:
|
|
seeds, compacted = unique_and_compact(nodes)
|
|
nodes_timestamp = None
|
|
(
|
|
compacted_node_pairs,
|
|
compacted_negative_srcs,
|
|
compacted_negative_dsts,
|
|
) = ({}, {}, {})
|
|
# Map back in same order as collect.
|
|
for etype, _ in node_pairs.items():
|
|
src_type, _, dst_type = etype_str_to_tuple(etype)
|
|
src = compacted[src_type].pop(0)
|
|
dst = compacted[dst_type].pop(0)
|
|
compacted_node_pairs[etype] = (src, dst)
|
|
if has_neg_src:
|
|
for etype, _ in neg_src.items():
|
|
src_type, _, _ = etype_str_to_tuple(etype)
|
|
compacted_negative_srcs[etype] = compacted[src_type].pop(0)
|
|
compacted_negative_srcs[etype] = compacted_negative_srcs[
|
|
etype
|
|
].view(neg_src[etype].shape)
|
|
if has_neg_dst:
|
|
for etype, _ in neg_dst.items():
|
|
_, _, dst_type = etype_str_to_tuple(etype)
|
|
compacted_negative_dsts[etype] = compacted[dst_type].pop(0)
|
|
compacted_negative_dsts[etype] = compacted_negative_dsts[
|
|
etype
|
|
].view(neg_dst[etype].shape)
|
|
else:
|
|
# Collect nodes from all types of input.
|
|
nodes = list(node_pairs)
|
|
nodes_timestamp = None
|
|
if use_timestamp:
|
|
# Timestamp for source and destination nodes are the same.
|
|
nodes_timestamp = [minibatch.timestamp, minibatch.timestamp]
|
|
if has_neg_src:
|
|
nodes.append(neg_src.view(-1))
|
|
if use_timestamp:
|
|
nodes_timestamp.append(
|
|
minibatch.timestamp.repeat_interleave(neg_src.shape[-1])
|
|
)
|
|
if has_neg_dst:
|
|
nodes.append(neg_dst.view(-1))
|
|
if use_timestamp:
|
|
nodes_timestamp.append(
|
|
minibatch.timestamp.repeat_interleave(neg_dst.shape[-1])
|
|
)
|
|
# Unique and compact the collected nodes.
|
|
if use_timestamp:
|
|
seeds, nodes_timestamp, compacted = compact_temporal_nodes(
|
|
nodes, nodes_timestamp
|
|
)
|
|
else:
|
|
seeds, compacted = unique_and_compact(nodes)
|
|
nodes_timestamp = None
|
|
# Map back in same order as collect.
|
|
compacted_node_pairs = tuple(compacted[:2])
|
|
compacted = compacted[2:]
|
|
if has_neg_src:
|
|
compacted_negative_srcs = compacted.pop(0)
|
|
# Since we need to calculate the neg_ratio according to the
|
|
# compacted_negatvie_srcs shape, we need to reshape it back.
|
|
compacted_negative_srcs = compacted_negative_srcs.view(
|
|
neg_src.shape
|
|
)
|
|
if has_neg_dst:
|
|
compacted_negative_dsts = compacted.pop(0)
|
|
# Same as above.
|
|
compacted_negative_dsts = compacted_negative_dsts.view(
|
|
neg_dst.shape
|
|
)
|
|
return (
|
|
seeds,
|
|
nodes_timestamp,
|
|
compacted_node_pairs,
|
|
compacted_negative_srcs if has_neg_src else None,
|
|
compacted_negative_dsts if has_neg_dst else None,
|
|
)
|
|
|
|
def _sample(self, minibatch):
|
|
(
|
|
minibatch.input_nodes,
|
|
minibatch.sampled_subgraphs,
|
|
) = self.sample_subgraphs(
|
|
minibatch._seed_nodes, minibatch._seeds_timestamp
|
|
)
|
|
return minibatch
|
|
|
|
def sampling_stages(self, datapipe):
|
|
"""The sampling stages are defined here by chaining to the datapipe. The
|
|
default implementation expects :meth:`sample_subgraphs` to be
|
|
implemented. To define fine-grained stages, this method should be
|
|
overridden.
|
|
"""
|
|
return datapipe.transform(self._sample)
|
|
|
|
@staticmethod
|
|
def _seeds_preprocess(minibatch):
|
|
"""Preprocess `seeds` in a minibatch to construct `unique_seeds`,
|
|
`node_timestamp` and `compacted_seeds` for further sampling. It
|
|
optionally incorporates timestamps for temporal graphs, organizing and
|
|
compacting seeds based on their types and timestamps.
|
|
|
|
Parameters
|
|
----------
|
|
minibatch: MiniBatch
|
|
The minibatch.
|
|
|
|
Returns
|
|
-------
|
|
unique_seeds: torch.Tensor or Dict[str, torch.Tensor]
|
|
A tensor or a dictionary of tensors representing the unique seeds.
|
|
In heterogeneous graphs, seeds are returned for each node type.
|
|
nodes_timestamp: None or a torch.Tensor or Dict[str, torch.Tensor]
|
|
Containing timestamps for each seed. This is only returned if
|
|
`minibatch` includes timestamps and the graph is temporal.
|
|
compacted_seeds: torch.tensor or a Dict[str, torch.Tensor]
|
|
Representation of compacted seeds corresponding to 'seeds', where
|
|
all node ids inside are compacted.
|
|
"""
|
|
use_timestamp = hasattr(minibatch, "timestamp")
|
|
seeds = minibatch.seeds
|
|
is_heterogeneous = isinstance(seeds, Dict)
|
|
if is_heterogeneous:
|
|
# Collect nodes from all types of input.
|
|
nodes = defaultdict(list)
|
|
nodes_timestamp = None
|
|
if use_timestamp:
|
|
nodes_timestamp = defaultdict(list)
|
|
for etype, pair in seeds.items():
|
|
assert pair.ndim == 1 or (
|
|
pair.ndim == 2 and pair.shape[1] == 2
|
|
), (
|
|
"Only tensor with shape 1*N and N*2 is "
|
|
+ f"supported now, but got {pair.shape}."
|
|
)
|
|
ntypes = etype[:].split(":")[::2]
|
|
pair = pair.view(pair.shape[0], -1)
|
|
if use_timestamp:
|
|
negative_ratio = (
|
|
pair.shape[0] // minibatch.timestamp[etype].shape[0] - 1
|
|
)
|
|
neg_timestamp = minibatch.timestamp[
|
|
etype
|
|
].repeat_interleave(negative_ratio)
|
|
for i, ntype in enumerate(ntypes):
|
|
nodes[ntype].append(pair[:, i])
|
|
if use_timestamp:
|
|
nodes_timestamp[ntype].append(
|
|
minibatch.timestamp[etype]
|
|
)
|
|
nodes_timestamp[ntype].append(neg_timestamp)
|
|
# Unique and compact the collected nodes.
|
|
if use_timestamp:
|
|
(
|
|
unique_seeds,
|
|
nodes_timestamp,
|
|
compacted,
|
|
) = compact_temporal_nodes(nodes, nodes_timestamp)
|
|
else:
|
|
unique_seeds, compacted = unique_and_compact(nodes)
|
|
nodes_timestamp = None
|
|
compacted_seeds = {}
|
|
# Map back in same order as collect.
|
|
for etype, pair in seeds.items():
|
|
if pair.ndim == 1:
|
|
compacted_seeds[etype] = compacted[etype].pop(0)
|
|
else:
|
|
src_type, _, dst_type = etype_str_to_tuple(etype)
|
|
src = compacted[src_type].pop(0)
|
|
dst = compacted[dst_type].pop(0)
|
|
compacted_seeds[etype] = torch.cat((src, dst)).view(2, -1).T
|
|
else:
|
|
# Collect nodes from all types of input.
|
|
nodes = [seeds.view(-1)]
|
|
nodes_timestamp = None
|
|
if use_timestamp:
|
|
# Timestamp for source and destination nodes are the same.
|
|
negative_ratio = (
|
|
seeds.shape[0] // minibatch.timestamp.shape[0] - 1
|
|
)
|
|
neg_timestamp = minibatch.timestamp.repeat_interleave(
|
|
negative_ratio
|
|
)
|
|
seeds_timestamp = torch.cat(
|
|
(minibatch.timestamp, neg_timestamp)
|
|
)
|
|
nodes_timestamp = [seeds_timestamp for _ in range(seeds.ndim)]
|
|
# Unique and compact the collected nodes.
|
|
if use_timestamp:
|
|
(
|
|
unique_seeds,
|
|
nodes_timestamp,
|
|
compacted,
|
|
) = compact_temporal_nodes(nodes, nodes_timestamp)
|
|
else:
|
|
unique_seeds, compacted = unique_and_compact(nodes)
|
|
nodes_timestamp = None
|
|
# Map back in same order as collect.
|
|
compacted_seeds = compacted[0].view(seeds.shape)
|
|
return (
|
|
unique_seeds,
|
|
nodes_timestamp,
|
|
compacted_seeds,
|
|
)
|
|
|
|
def sample_subgraphs(self, seeds, seeds_timestamp):
|
|
"""Sample subgraphs from the given seeds, possibly with temporal constraints.
|
|
|
|
Any subclass of SubgraphSampler should implement this method.
|
|
|
|
Parameters
|
|
----------
|
|
seeds : Union[torch.Tensor, Dict[str, torch.Tensor]]
|
|
The seed nodes.
|
|
|
|
seeds_timestamp : Union[torch.Tensor, Dict[str, torch.Tensor]]
|
|
The timestamps of the seed nodes. If given, the sampled subgraphs
|
|
should not contain any nodes or edges that are newer than the
|
|
timestamps of the seed nodes. Default: None.
|
|
|
|
Returns
|
|
-------
|
|
Union[torch.Tensor, Dict[str, torch.Tensor]]
|
|
The input nodes.
|
|
List[SampledSubgraph]
|
|
The sampled subgraphs.
|
|
|
|
Examples
|
|
--------
|
|
>>> @functional_datapipe("my_sample_subgraph")
|
|
>>> class MySubgraphSampler(SubgraphSampler):
|
|
>>> def __init__(self, datapipe, graph, fanouts):
|
|
>>> super().__init__(datapipe)
|
|
>>> self.graph = graph
|
|
>>> self.fanouts = fanouts
|
|
>>> def sample_subgraphs(self, seeds):
|
|
>>> # Sample subgraphs from the given seeds.
|
|
>>> subgraphs = []
|
|
>>> subgraphs_nodes = []
|
|
>>> for fanout in reversed(self.fanouts):
|
|
>>> subgraph = self.graph.sample_neighbors(seeds, fanout)
|
|
>>> subgraphs.insert(0, subgraph)
|
|
>>> subgraphs_nodes.append(subgraph.nodes)
|
|
>>> seeds = subgraph.nodes
|
|
>>> subgraphs_nodes = torch.unique(torch.cat(subgraphs_nodes))
|
|
>>> return subgraphs_nodes, subgraphs
|
|
"""
|
|
raise NotImplementedError
|