LayerNeighborSampler

class dgl.graphbolt.LayerNeighborSampler(datapipe, graph, fanouts, replace=False, prob_name=None, deduplicate=True, layer_dependency=False, batch_dependency=1, overlap_fetch=False, num_gpu_cached_edges=0, gpu_cache_threshold=1, cooperative=False, asynchronous=False)[source]

基类:NeighborSamplerImpl

从图中采样层邻居边并返回子图。

函数名: sample_layer_neighbor.

此采样器通过劳动采样构建多层 GNN 的节点表示计算依赖关系,劳动采样方法来自 NeurIPS 2023 论文 Layer-Neighbor Sampling – Defusing Neighborhood Explosion in GNNs

Layer-Neighbor 采样器负责从给定数据中采样一个子图。它返回一个诱导子图以及压缩信息。在节点分类任务的上下文中,邻居采样器直接利用提供的节点作为种子节点。然而,在涉及链接预测的场景中,该过程需要另一个预处理操作。也就是说,从给定的节点对(包括正负节点对)中收集唯一节点,并将这些节点用作后续步骤的种子节点。当图是异构的时,minibatch 中的采样子图将包含所有边类型,即使采样后它们为空。

实现了论文附录 A.3 中描述的方法。与 dgl.dataloading.LaborSampler 类似,但它使用序列泊松采样代替泊松采样,以使每个顶点的采样边计数像 NeighborSampler 一样是确定的。因此,它是 NeighborSampler 的直接替代品。然而,与 NeighborSampler 不同的是,它在多层 GNN 场景中采样更少的顶点和边,而不会损害训练迭代的收敛速度。

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

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

  • fanouts (list[torch.Tensor]) – 每个节点需要采样的边数,可以考虑或不考虑边类型。此参数的长度隐式表示正在进行的采样层数。

  • replace (bool) – 布尔值,指示采样是带替换还是不带替换。如果为 True,则一个值可以被多次选中。否则,每个值只能被选中一次。

  • prob_name (str, optional) – 用作每个节点采样权重的边属性名称。此属性张量应包含对应于节点每个邻居边的(未归一化)概率。它必须是 1D 浮点或布尔张量,元素数量等于总边数。

  • deduplicate (bool) – 布尔值,指示是否对跳之间的种子进行去重。如果为 True,则种子中的相同元素将仅保留一个。否则,将保留所有相同元素。

  • layer_dependency (bool) – 布尔值,指示不同层是否应使用相同的随机变量。这会减少采样节点的数量,并将 LayerNeighborSampler 转换为一种子图采样方法。后续层将保证采样到与前一层重叠的邻居。

  • batch_dependency (int) – 指定连续的 minibatch 是否应使用相似的随机变量。这会提高采样节点和边的时序访问局部性。将其设置为 \(\kappa\) 会使随机变量的变化减缓,减缓比例为 \(\frac{1}{\kappa}\)。实现了 arXiv:2310.12403 中的依赖式 minibatching 方法。

  • overlap_fetch (bool, optional) – 如果为 True,则数据加载器将使用备用 CUDA 流将 UVA 图提取操作与其余操作重叠。如果您已将图移至固定内存以获得最佳性能,则应启用此选项。默认为 False。

  • num_gpu_cached_edges (int, optional) – 如果为正且 overlap_graph_fetch 为 True,则 GPU 将缓存频繁访问的顶点邻域,以减少由于固定图访问引起的 PCI-e 带宽需求。

  • gpu_cache_threshold (int, optional) – 确定顶点需要被访问多少次后,其邻域才会被缓存到 GPU 上。

  • cooperative (bool, optional) – 布尔值,指示是否启用协作式 Minibatching,该方法最初在 `Deep Graph Library PR#4337<https://github.com/dmlc/dgl/pull/4337>`__ 中提出,后来在 Cooperative Minibatching in Graph Neural Networks 中首次完整描述。GPU 之间的协作消除了由于种子节点在执行 GNN Minibatching 时采样的 k 跳邻域重叠而在 GPU 之间执行的重复工作。

  • asynchronous (bool) – 布尔值,指示采样和压缩阶段是否应在后台线程中运行,以隐藏 CPU GPU 同步的延迟。仅在 GPU 上采样时应启用此选项。

示例

>>> import dgl.graphbolt as gb
>>> import torch
>>> indptr = torch.LongTensor([0, 2, 4, 5, 6, 7 ,8])
>>> indices = torch.LongTensor([1, 2, 0, 3, 5, 4, 3, 5])
>>> graph = gb.fused_csc_sampling_graph(indptr, indices)
>>> seeds = torch.LongTensor([[0, 1], [1, 2]])
>>> item_set = gb.ItemSet(seeds, names="seeds")
>>> item_sampler = gb.ItemSampler(item_set, batch_size=1,)
>>> neg_sampler = gb.UniformNegativeSampler(item_sampler, graph, 2)
>>> fanouts = [torch.LongTensor([5]),
...     torch.LongTensor([10]),torch.LongTensor([15])]
>>> subgraph_sampler = gb.LayerNeighborSampler(neg_sampler, graph, fanouts)
>>> next(iter(subgraph_sampler)).sampled_subgraphs
[SampledSubgraphImpl(sampled_csc=CSCFormatBase(
        indptr=tensor([0, 2, 4, 5, 6, 7, 8]),
        indices=tensor([1, 3, 0, 4, 2, 2, 5, 4]),
    ),
    original_row_node_ids=tensor([0, 1, 5, 2, 3, 4]),
    original_edge_ids=None,
    original_column_node_ids=tensor([0, 1, 5, 2, 3, 4]),
),
SampledSubgraphImpl(sampled_csc=CSCFormatBase(
        indptr=tensor([0, 2, 4, 5, 6, 7]),
        indices=tensor([1, 3, 0, 4, 2, 2, 5]),
    ),
    original_row_node_ids=tensor([0, 1, 5, 2, 3, 4]),
    original_edge_ids=None,
    original_column_node_ids=tensor([0, 1, 5, 2, 3]),
),
SampledSubgraphImpl(sampled_csc=CSCFormatBase(
        indptr=tensor([0, 2, 4, 5, 6]),
        indices=tensor([1, 3, 0, 4, 2, 2]),
    ),
    original_row_node_ids=tensor([0, 1, 5, 2, 3]),
    original_edge_ids=None,
    original_column_node_ids=tensor([0, 1, 5, 2]),
)]
>>> next(iter(subgraph_sampler)).compacted_seeds
tensor([[0, 1], [0, 2], [0, 3]])
>>> next(iter(subgraph_sampler)).labels
tensor([1., 0., 0.])
>>> next(iter(subgraph_sampler)).indexes
tensor([0, 0, 0])