dgl.sampling.sample_labors

dgl.sampling.sample_labors(g, nodes, fanout, edge_dir='in', prob=None, importance_sampling=0, random_seed=None, seed2_contribution=0, copy_ndata=True, copy_edata=True, exclude_edges=None, output_device=None)[source]

该采样器通过劳动采样（labor sampling）构建多层 GNN 节点表示的计算依赖关系，该算法出自 NeurIPS 2023 的论文 Layer-Neighbor Sampling – Defusing Neighborhood Explosion in GNNs

该采样器将使每个节点针对每种边类型从固定数量的邻居那里收集消息。默认参数下，邻居是均匀采样的。对于每个将被考虑采样的顶点 t，都会有一个单一的随机变量 r_t。

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

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

参数:

  • g (DGLGraph) – 待采样的图，允许包含多种节点或边类型。可以在 CPU 或 GPU 上。

  • nodes (tensor 或 dict) –

    要从中采样邻居的节点 ID。

    此参数可以接受单个 ID 张量，或一个包含节点类型和 ID 张量的字典。如果给定单个张量，则图必须只有一种节点类型。

  • fanout (int 或 dict[etype, int]) –

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

    此参数可以接受单个 int，或一个包含边类型和 int 的字典。如果给定单个 int，DGL 将为每个节点每种边类型采样此数量的边。

    如果对于单一边类型给定 -1，则将选择该边类型的所有邻居边。

  • edge_dir (str, 可选) –

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

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

  • prob (str, 可选) –

    用作与节点每个邻居边相关的（未归一化）概率的特征名称。该特征对于每条边必须只有一个元素。

    特征必须是非负浮点数，并且每个节点的入边/出边特征的总和必须为正（尽管它们不必总和为一）。否则，结果将是未定义的。

    如果 prob 不是 None，则不支持 GPU 采样。

  • importance_sampling (int, 可选) – 是否使用重要性采样或均匀采样，使用负值会优化重要性采样概率直到收敛，而使用正值会运行多次优化步骤。如果值为 i，则使用 LABOR-i 变体。

  • random_seed (tensor) –

    一个包含一个元素的 int64 张量。

    传入的 random_seed 使得对于任何种子顶点 s 及其邻居 t，在调用此函数时使用相同的 random_seed 时，生成的随机变量 r_t 是相同的。当作为同一批次的一部分进行采样时，需要相同的种子，以便 LABOR 可以进行全局采样。一个例子是对于异构图，为每种边类型传递一个相同的 random_seed。与为每种边类型使用不同的 random_seed 相比，这将采样少得多的顶点。如果在异构图中为每种边类型单独调用此函数并使用不同的 random_seed，那么它将在每种边类型上局部运行 LABOR，导致采样更多顶点。

    如果在调用此函数时没有传入 random_seed，DGL 将生成一个随机数作为 random_seed。如果将此函数的多次调用作为单个批次的一部分进行采样，请使用相同的 random_seed。

  • seed2_contribution (float, 可选) – 一个介于 [0, 1) 之间的浮点值，用于确定第二个随机种子在生成 LABOR 采样算法的随机变量时的贡献度。

  • copy_ndata (bool, 可选) –

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

    (默认值: True)

  • copy_edata (bool, 可选) –

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

    (默认值: True)

  • exclude_edges (tensor 或 dict) –

    在为种子节点采样邻居时要排除的边 ID。

    此参数可以接受单个 ID 张量，或一个包含边类型和 ID 张量的字典。如果给定单个张量，则图必须只有一种节点类型。

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

返回值:

一个采样得到的子图，仅包含采样的邻居边以及边权重。

返回值类型:

tuple(DGLGraph, list[Tensor])

注意事项

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

示例

假设您有以下图

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

以及权重

>>> g.edata['prob'] = torch.FloatTensor([0., 1., 0., 1., 0., 1.])

为节点 0 和节点 1 采样一条入边

>>> sg = dgl.sampling.sample_labors(g, [0, 1], 1)
>>> sg.edges(order='eid')
(tensor([1, 0]), tensor([0, 1]))
>>> sg.edata[dgl.EID]
tensor([2, 0])

使用边特征 prob 中的概率为节点 0 和节点 1 采样一条入边

>>> sg = dgl.sampling.sample_labors(g, [0, 1], 1, prob='prob')
>>> sg.edges(order='eid')
(tensor([2, 1]), tensor([0, 1]))

当 fanout 大于实际邻居数量且不允许重复采样时，DGL 将改为选择所有邻居

>>> sg = dgl.sampling.sample_labors(g, [0, 1], 3)
>>> sg.edges(order='eid')
(tensor([1, 2, 0, 1]), tensor([0, 0, 1, 1]))

在为种子节点采样时排除某些 EID

>>> g = dgl.graph(([0, 0, 1, 1, 2, 2], [1, 2, 0, 1, 2, 0]))
>>> g_edges = g.all_edges(form='all')``
(tensor([0, 0, 1, 1, 2, 2]), tensor([1, 2, 0, 1, 2, 0]), tensor([0, 1, 2, 3, 4, 5]))
>>> sg = dgl.sampling.sample_labors(g, [0, 1], 3, exclude_edges=[0, 1, 2])
>>> sg.all_edges(form='all')
(tensor([2, 1]), tensor([0, 1]), tensor([0, 1]))
>>> sg.has_edges_between(g_edges[0][:3],g_edges[1][:3])
tensor([False, False, False])
>>> g = dgl.heterograph({
...   ('drug', 'interacts', 'drug'): ([0, 0, 1, 1, 3, 2], [1, 2, 0, 1, 2, 0]),
...   ('drug', 'interacts', 'gene'): ([0, 0, 1, 1, 2, 2], [1, 2, 0, 1, 2, 0]),
...   ('drug', 'treats', 'disease'): ([0, 0, 1, 1, 2, 2], [1, 2, 0, 1, 2, 0])})
>>> g_edges = g.all_edges(form='all', etype=('drug', 'interacts', 'drug'))
(tensor([0, 0, 1, 1, 3, 2]), tensor([1, 2, 0, 1, 2, 0]), tensor([0, 1, 2, 3, 4, 5]))
>>> excluded_edges  = {('drug', 'interacts', 'drug'): g_edges[2][:3]}
>>> sg = dgl.sampling.sample_labors(g, {'drug':[0, 1]}, 3, exclude_edges=excluded_edges)
>>> sg.all_edges(form='all', etype=('drug', 'interacts', 'drug'))
(tensor([2, 1]), tensor([0, 1]), tensor([0, 1]))
>>> sg.has_edges_between(g_edges[0][:3],g_edges[1][:3],etype=('drug', 'interacts', 'drug'))
tensor([False, False, False])