NeighborSampler

class dgl.graphbolt.NeighborSampler(datapipe, graph, fanouts, replace=False, prob_name=None, deduplicate=True, overlap_fetch=False, num_gpu_cached_edges=0, gpu_cache_threshold=1, cooperative=False, asynchronous=False)[源代码]

基础类: NeighborSamplerImpl

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

函数名: sample_neighbor。

邻居采样器负责从给定数据中采样一个子图。它返回一个诱导子图以及紧凑的信息。在节点分类任务中，邻居采样器直接使用提供的节点作为种子节点。然而，在涉及链接预测的场景中，该过程需要另一个预处理操作。也就是说，从给定的节点对（包括正负节点对）中收集唯一的节点，并将这些节点用作后续步骤的种子节点。当图是异构的（hetero）时，小批量中采样的子图将包含每种边类型，即使采样后它是空的。

参数：

datapipe (DataPipe) – 数据管道。
graph (FusedCSCSamplingGraph) – 执行子图采样的图。
fanouts (list[torch.Tensor] or list[int]) – 对于每个节点，采样边的数量，可以考虑或不考虑边类型。此参数的长度隐式表示正在进行的采样层数。注意：fanout 顺序是从最外层到最内层。例如，fanout ‘[15, 10, 5]’ 表示最外层采样 15 条边，中间层采样 10 条边，最内层采样 5 条边。
replace (bool) – 布尔值，指示采样是有放回还是无放回。如果为 True，则一个值可以多次被选中。否则，每个值只能被选中一次。
prob_name (str, optional) – 用于作为每个节点采样权重的边属性名称。此属性张量应包含与节点每个邻居边对应的（未归一化的）概率。它必须是一个 1D 浮点或布尔张量，元素数量等于边的总数。
deduplicate (bool) – 布尔值，指示层之间的种子是否去重。如果为 True，种子中的相同元素将被删除，只保留一个。否则，相同元素将保留。
overlap_fetch (bool, optional) – 如果为 True，数据加载器将使用另一个 CUDA 流将 UVA 图获取操作与其余操作重叠进行。如果您已将图移动到锁定内存（pinned memory）以获得最佳性能，应启用此选项。默认为 False。
num_gpu_cached_edges (int, optional) – 如果为正且 overlap_graph_fetch 为 True，则 GPU 将缓存频繁访问的顶点邻域，以减少因访问锁定内存图而导致的 PCI-e 带宽需求。
gpu_cache_threshold (int, optional) – 确定一个顶点需要被访问多少次，其邻域才会被缓存到 GPU 上。
cooperative (bool, optional) – 布尔值，指示是否启用协作式小批量（Cooperative Minibatching），该方法最初在 `Deep Graph Library PR#4337<https://github.com/dmlc/dgl/pull/4337>`__ 中提出，后来在 Graph Neural Networks 中的协作式小批量中首次完整描述。GPU 之间的协作消除了由于在执行 GNN 小批量处理时种子节点采样的 k 跳邻域重叠而在 GPU 之间执行的重复工作。
asynchronous (bool) – 布尔值，指示采样和压缩阶段是否应在后台线程中运行，以隐藏 CPU GPU 同步的延迟。仅在 GPU 上采样时应启用此选项。

示例

>>> import torch
>>> import dgl.graphbolt as gb
>>> 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")
>>> datapipe = gb.ItemSampler(item_set, batch_size=1)
>>> datapipe = datapipe.sample_uniform_negative(graph, 2)
>>> datapipe = datapipe.sample_neighbor(graph, [5, 10, 15])
>>> next(iter(datapipe)).sampled_subgraphs
[SampledSubgraphImpl(sampled_csc=CSCFormatBase(
        indptr=tensor([0, 2, 4, 5, 6, 7, 8]),
        indices=tensor([1, 4, 0, 5, 5, 3, 3, 2]),
    ),
    original_row_node_ids=tensor([0, 1, 4, 5, 2, 3]),
    original_edge_ids=None,
    original_column_node_ids=tensor([0, 1, 4, 5, 2, 3]),
),
SampledSubgraphImpl(sampled_csc=CSCFormatBase(
        indptr=tensor([0, 2, 4, 5, 6, 7, 8]),
        indices=tensor([1, 4, 0, 5, 5, 3, 3, 2]),
    ),
    original_row_node_ids=tensor([0, 1, 4, 5, 2, 3]),
    original_edge_ids=None,
    original_column_node_ids=tensor([0, 1, 4, 5, 2, 3]),
),
SampledSubgraphImpl(sampled_csc=CSCFormatBase(
        indptr=tensor([0, 2, 4, 5, 6]),
        indices=tensor([1, 4, 0, 5, 5, 3]),
    ),
    original_row_node_ids=tensor([0, 1, 4, 5, 2, 3]),
    original_edge_ids=None,
    original_column_node_ids=tensor([0, 1, 4, 5]),
)]