dgl.sampling.select_topk

dgl.sampling.select_topk(g, k, weight, nodes=None, edge_dir='in', ascending=False, copy_ndata=True, copy_edata=True, output_device=None)[source]

选择给定节点的k个最大（或最小）权重的相邻边，并返回诱导子图。

对于每个节点，将选择具有最大（或当ascending == True时为最小）权重的若干入站（或当edge_dir == 'out'时为出站）边。 返回的图将包含原始图中的所有节点，但仅包含采样的边。

Node/edge features are not preserved. The original IDs of the sampled edges are stored as the dgl.EID feature in the returned graph.

Parameters:

  • g (DGLGraph) – 图。必须在CPU上。

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

    每个节点在每个边类型上要选择的边的数量。

    这个参数可以接受一个整数或一个边类型和整数的字典。 如果给定一个整数，DGL将为每个节点为每个边类型选择这个数量的边。

    如果为单个边类型给定-1，将选择该边类型的所有相邻边。

  • weight (str) – 与每条边关联的权重特征名称。每条边的特征应只有一个元素。特征可以是 int32/64 或 float32/64。

  • nodes (tensor 或 dict, 可选) –

    从中采样邻居的节点ID。

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

    如果为None，DGL将为所有节点选择边。

  • edge_dir (str, optional) –

    Determines whether to sample inbound or outbound edges.

    Can take either in for inbound edges or out for outbound edges.

  • ascending (bool, optional) – 如果为True，DGL将返回具有k个最小权重的边，而不是k个最大权重的边。

  • copy_ndata (bool, optional) –

    If True, the node features of the new graph are copied from the original graph. If False, the new graph will not have any node features.

    (Default: True)

  • copy_edata (bool, optional) –

    If True, the edge features of the new graph are copied from the original graph. If False, the new graph will not have any edge features.

    (Default: True)

  • output_device (Framework-specific device context object, optional) – The output device. Default is the same as the input graph.

Returns:

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

Return type:

DGLGraph

注释

If copy_ndata or copy_edata is True, same tensors are used as the node or edge features of the original graph and the new graph. As a result, users should avoid performing in-place operations on the node features of the new graph to avoid feature corruption.

示例

>>> g = dgl.graph(([0, 0, 1, 1, 2, 2], [1, 2, 0, 1, 2, 0]))
>>> g.edata['weight'] = torch.FloatTensor([0, 1, 0, 1, 0, 1])
>>> sg = dgl.sampling.select_topk(g, 1, 'weight')
>>> sg.edges(order='eid')
(tensor([2, 1, 0]), tensor([0, 1, 2]))