InSubgraphSampler

class dgl.graphbolt.InSubgraphSampler(datapipe, graph)[源码]

基类: SubgraphSampler

采样给定节点的入边诱导的子图。

函数名: sample_in_subgraph.

入边子图采样器负责从给定数据中采样一个子图,并返回诱导子图及压缩信息。

参数:
  • datapipe (DataPipe) – 数据管道。

  • graph (FusedCSCSamplingGraph) – 执行入边子图采样的图。

示例

>>> import dgl.graphbolt as gb
>>> import torch
>>> indptr = torch.LongTensor([0, 3, 5, 7, 9, 12, 14])
>>> indices = torch.LongTensor([0, 1, 4, 2, 3, 0, 5, 1, 2, 0, 3, 5, 1, 4])
>>> graph = gb.fused_csc_sampling_graph(indptr, indices)
>>> item_set = gb.ItemSet(len(indptr) - 1, names="seeds")
>>> item_sampler = gb.ItemSampler(item_set, batch_size=2)
>>> insubgraph_sampler = gb.InSubgraphSampler(item_sampler, graph)
>>> for _, data in enumerate(insubgraph_sampler):
...     print(data.sampled_subgraphs[0].sampled_csc)
...     print(data.sampled_subgraphs[0].original_row_node_ids)
...     print(data.sampled_subgraphs[0].original_column_node_ids)
CSCFormatBase(indptr=tensor([0, 3, 5]),
            indices=tensor([0, 1, 2, 3, 4]),
)
tensor([0, 1, 4, 2, 3])
tensor([0, 1])
CSCFormatBase(indptr=tensor([0, 2, 4]),
            indices=tensor([2, 3, 4, 0]),
)
tensor([2, 3, 0, 5, 1])
tensor([2, 3])
CSCFormatBase(indptr=tensor([0, 3, 5]),
            indices=tensor([2, 3, 1, 4, 0]),
)
tensor([4, 5, 0, 3, 1])
tensor([4, 5])
sample_subgraphs(seeds, seeds_timestamp, seeds_pre_time_window=None)[源码]

从给定种子中采样子图,可能包含时间约束。

SubgraphSampler 的任何子类都应实现此方法。

参数:
  • seeds (Union[torch.Tensor, Dict[str, torch.Tensor]]) – 种子节点。

  • seeds_timestamp (Union[torch.Tensor, Dict[str, torch.Tensor]]) – 种子节点的时间戳。如果给定,采样的子图不应包含任何比种子节点时间戳更新的节点或边。默认值:None。

  • seeds_pre_time_window (Union[torch.Tensor, Dict[str, torch.Tensor]]) – 节点的时间窗口表示 seeds_timestamp 之前的一段时间。如果提供,将仅过滤时间戳在 [seeds_timestamp - seeds_pre_time_window, seeds_timestamp] 范围内的邻居和相关边。

返回值:

  • Union[torch.Tensor, Dict[str, torch.Tensor]] – 输入节点。

  • List[SampledSubgraph] – 采样的子图。

示例

>>> @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