dgl.topk_nodesΒΆ

dgl.topk_nodes(graph, feat, k, *, descending=True, sortby=None, ntype=None)[source]ΒΆ

Return a graph-level representation by a graph-wise top-k on node features feat in graph by feature at index sortby.

If descending is set to False, return the k smallest elements instead.

If sortby is set to None, the function would perform top-k on all dimensions independently, equivalent to calling torch.topk(graph.ndata[feat], dim=0).

Parameters

  • graph (DGLGraph) – The graph.

  • feat (str) – The feature field.

  • k (int) – The k in β€œtop-k”

  • descending (bool) – Controls whether to return the largest or smallest elements.

  • sortby (int, optional) – Sort according to which feature. If is None, all features are sorted independently.

  • ntype (str, optional) – Node type. Can be omitted if there is only one node type in the graph.

Returns

  • sorted_feat (Tensor) – A tensor with shape \((B, K, D)\), where \(B\) is the batch size of the input graph.

  • sorted_idx (Tensor) – A tensor with shape \((B, K)\) if sortby is set to None), where \(B\) is the batch size of the input graph, \(D\) is the feature size.

Notes

If an example has \(n\) nodes and \(n<k\), the sorted_feat tensor will pad the \(n+1\) to \(k\) th rows with zero;

Examples

>>> import dgl
>>> import torch as th

Create two DGLGraph objects and initialize their node features.

>>> g1 = dgl.graph(([0, 1], [2, 3]))              # Graph 1
>>> g1.ndata['h'] = th.rand(4, 5)
>>> g1.ndata['h']
tensor([[0.0297, 0.8307, 0.9140, 0.6702, 0.3346],
        [0.5901, 0.3030, 0.9280, 0.6893, 0.7997],
        [0.0880, 0.6515, 0.4451, 0.7507, 0.5297],
        [0.5171, 0.6379, 0.2695, 0.8954, 0.5197]])

>>> g2 = dgl.graph(([0, 1, 2], [2, 3, 4]))       # Graph 2
>>> g2.ndata['h'] = th.rand(5, 5)
>>> g2.ndata['h']
tensor([[0.3168, 0.3174, 0.5303, 0.0804, 0.3808],
        [0.1323, 0.2766, 0.4318, 0.6114, 0.1458],
        [0.1752, 0.9105, 0.5692, 0.8489, 0.0539],
        [0.1931, 0.4954, 0.3455, 0.3934, 0.0857],
        [0.5065, 0.5182, 0.5418, 0.1520, 0.3872]])

Top-k over node attribute h in a batched graph.

>>> bg = dgl.batch([g1, g2], ndata=['h'])
>>> dgl.topk_nodes(bg, 'h', 3)
(tensor([[[0.5901, 0.8307, 0.9280, 0.8954, 0.7997],
          [0.5171, 0.6515, 0.9140, 0.7507, 0.5297],
          [0.0880, 0.6379, 0.4451, 0.6893, 0.5197]],
         [[0.5065, 0.9105, 0.5692, 0.8489, 0.3872],
          [0.3168, 0.5182, 0.5418, 0.6114, 0.3808],
          [0.1931, 0.4954, 0.5303, 0.3934, 0.1458]]]), tensor([[[1, 0, 1, 3, 1],
          [3, 2, 0, 2, 2],
          [2, 3, 2, 1, 3]],
         [[4, 2, 2, 2, 4],
          [0, 4, 4, 1, 0],
          [3, 3, 0, 3, 1]]]))

Top-k over node attribute h along the last dimension in a batched graph. (used in SortPooling)

>>> dgl.topk_nodes(bg, 'h', 3, sortby=-1)
(tensor([[[0.5901, 0.3030, 0.9280, 0.6893, 0.7997],
          [0.0880, 0.6515, 0.4451, 0.7507, 0.5297],
          [0.5171, 0.6379, 0.2695, 0.8954, 0.5197]],
         [[0.5065, 0.5182, 0.5418, 0.1520, 0.3872],
          [0.3168, 0.3174, 0.5303, 0.0804, 0.3808],
          [0.1323, 0.2766, 0.4318, 0.6114, 0.1458]]]), tensor([[1, 2, 3],
         [4, 0, 1]]))

Top-k over node attribute h in a single graph.

>>> dgl.topk_nodes(g1, 'h', 3)
(tensor([[[0.5901, 0.8307, 0.9280, 0.8954, 0.7997],
          [0.5171, 0.6515, 0.9140, 0.7507, 0.5297],
          [0.0880, 0.6379, 0.4451, 0.6893, 0.5197]]]), tensor([[[1, 0, 1, 3, 1],
          [3, 2, 0, 2, 2],
          [2, 3, 2, 1, 3]]]))