dgl.sampling.sample_neighbors_biased

dgl.sampling.sample_neighbors_biased(g, nodes, fanout, bias, edge_dir='in', tag_offset_name='_TAG_OFFSET', replace=False, copy_ndata=True, copy_edata=True, output_device=None)[source]

对给定节点的邻居边进行采样并返回诱导子图，其中每个邻居被选中的概率由其标签决定。

对于每个节点，将随机选择一定数量的入边（或 edge_dir == 'out' 时的出边）。返回的图将包含原始图中的所有节点，但仅包含采样的边。

此版本的邻居采样支持相邻节点类型不同时具有不同采样概率的场景。每个节点被分配一个整数（称为 标签），它表示节点的类型。标签在同构图框架下是节点类型的一种模拟。具有相同标签的节点共享相同的概率。

例如，假设一个节点有 \(N+M\) 个邻居，其中 \(N\) 个的标签为 0，\(M\) 个的标签为 1。假设标签为 0 的节点被选中的非归一化概率为 \(p\)，而标签为 1 的节点被选中的非归一化概率为 \(q\)。此函数首先根据非归一化概率分布 \(\frac{P(tag=0)}{P(tag=1)}=\frac{Np}{Mq}\) 选择一个标签，然后进行均匀采样以获得具有所选标签的节点。

为了使采样更高效，输入图的 CSC 矩阵（如果 edge_dir='out' 则为 CSR 矩阵）必须按标签排序。API sort_csc_by_tag() 和 sort_csr_by_tag() 就是为此目的设计的，它们会在内部按标签重新排序邻居，以便具有相同标签的邻居存储在连续的范围内。这两个 API 还会将这些范围的偏移量存储在一个节点特征中，该特征的名称为 tag_offset_name。

请确保在调用此函数之前，图的 CSR（或 CSC）矩阵已排序。 此函数本身不会检查输入图是否已排序。请注意，输入的 tag_offset_name 应与排序函数中的一致。

仅支持同构图或二分图。对于二分图，当 edge_dir='in' 时，将使用源节点的标签偏移量进行采样（或当 edge_dir='out' 时使用目标节点的标签偏移量）。

节点/边特征不会被保留。采样边的原始 ID 作为 dgl.EID 特征存储在返回的图中。

参数:

  • g (DGLGraph) – 图。必须是同构图或二分图（只有一种边类型）。必须在 CPU 上。

  • nodes (tensor 或 list) – 要采样邻居的节点 ID。

  • fanout (int) –

    对于每种边类型的每个节点要采样的边数。

    如果给定 -1，将选择所有概率非零的邻居边。

  • bias (tensor 或 list) –

    与每个标签相关的（非归一化）概率。其长度应等于标签的数量。

    此数组的条目必须是非负浮点数。否则，结果将未定义。

  • edge_dir (str, 可选) –

    确定是采样入边还是出边。

    可以是 in 表示入边，也可以是 out 表示出边。

  • tag_offset_name (str, 可选) –

    存储标签偏移量的节点特征的名称。

    (默认值: “_TAG_OFFSET”)

  • replace (bool, 可选) – 如果为 True，则进行有放回采样。

  • copy_ndata (bool, 可选) –

    如果为 True，则新图的节点特征会从原始图复制。如果为 False，则新图将没有任何节点特征。

    (默认值: True)

  • copy_edata (bool, 可选) –

    如果为 True，则新图的边特征会从原始图复制。如果为 False，则新图将没有任何边特征。

    (默认值: True)

  • output_device (框架特定的设备上下文对象, 可选) – 输出设备。默认与输入图相同。

返回值:

一个采样的子图，仅包含采样的邻居边。它位于 CPU 上。

返回类型:

DGLGraph

注意

如果 copy_ndata 或 copy_edata 为 True，则新图的节点或边特征使用与原始图相同的张量。因此，用户应避免在新图的节点特征上执行就地操作，以避免特征损坏。

另请参阅

dgl.sort_csc_by_tag, dgl.sort_csr_by_tag

示例

假设您有以下图

>>> g = dgl.graph(([0, 0, 1, 1, 2, 2], [1, 2, 0, 1, 2, 0]))

以及标签

>>> tag = torch.IntTensor([0, 0, 1])

排序图（必要！）

>>> g_sorted = dgl.transforms.sort_csr_by_tag(g, tag)
>>> g_sorted.ndata['_TAG_OFFSET']
tensor([[0, 1, 2],
        [0, 2, 2],
        [0, 1, 2]])

设置每个标签的概率

>>> bias = torch.tensor([1.0, 0.001])
>>> # node 2 is almost impossible to be sampled because it has tag 1.

采样节点 0 和节点 2 的一条出边

>>> sg = dgl.sampling.sample_neighbors_biased(g_sorted, [0, 2], 1, bias, edge_dir='out')
>>> sg.edges(order='eid')
(tensor([0, 2]), tensor([1, 0]))
>>> sg.edata[dgl.EID]
tensor([0, 5])

当 fanout 大于实际邻居数且无放回采样时，DGL 将转而选择所有邻居

>>> sg = dgl.sampling.sample_neighbors_biased(g_sorted, [0, 2], 3, bias, edge_dir='out')
>>> sg.edges(order='eid')
(tensor([0, 0, 2, 2]), tensor([1, 2, 0, 2]))