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
ingraph
by feature at indexsortby
.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 callingtorch.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]]]))