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])