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]]]))