Sparse topk

[yorard123]

Hi,

I was wondering if is there an easy way to find the topk values of a prediction in pytorch geometric. Example imagine that you have the following data structure

data.x = [0.2, 0.1, 0.1, 0.4, 0.4 , 0.05, 0.05, 0.1, 0.05, 0,01, 0.2, 0.15, 0.05, 0.4]
data.batch = [0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1]

what I would like to do for instance is find the k top values independently for each sample. I would expected something like:
k = 2
data.topk = [ 0.4, 0.4, 0.4, 0.2]
data.batch = [0, 0, 1, 1]

Is there any method in pytorch geometric to do that? I have been checking the documentation but I was not able to find something like this.

EDIT:

I found that there is this method https://pytorch-geometric.readthedocs.io/en/latest/_modules/torch_geometric/nn/glob/sort.html#global_sort_pool that it should do the trick, however, it is not returning the location of the maximum values, and this is something that I should need to recover the new batch. Is there any method that I am missing?

Thanks!

[QuanticDisaster]

It is implemented in pytorch, I think this is what you are looking for : https://pytorch.org/docs/stable/generated/torch.topk.html

[yorard123]

This does not work with the pytorch geometric structure, where the batch is defined using an index vector.

[rusty1s]

There exists a topk function deeply hidden in PyG :)

An alternative approach is to utilize PyTorch's topk function (as @QuanticDisaster) on a dense batch:

x, mask = to_dense_batch(x, batch)
torch.topk(x, k=2, dim=1)

[yorard123]

Thanks!! Is it possible that this topk function does not appear to the documentation? I think it could be very useful to have it in there. I have another related question to this one. I would like to know if is there a sort function to tackle the same problem. The idea it would be to sort the values of x independently and have as output the perm vector.

[rusty1s]

Yes, we should make it a public function visible in the documentation. Do you have any interest in contributing this feature? :)
We can also add support for batch-wise sorting, but in theory, this should already be supported by topk with ratio=1.0.

[wsad1]

Note the below approach would only work if all the graphs have the same number of nodes. If that's not the case then for some graphs x would be filled by 0s which would end up giving you the wrong topk result.

x, mask = to_dense_batch(x, batch)
torch.topk(x, k=2, dim=1)

So making PyG's topk visible is a good idea.

[yorard123]

Do you have any interest in contributing this feature? :)
Do you mean about making visible the topk?

Note the below approach would only work if all the graphs have the same number of nodes. When the nodes does not have the same number of nodes is it true that it could give you wrong topk, but you can 'solved it' with the fill_value parameter of the to_dense_batch method choosing a very low value. However, I believe that making topk visible is good, in fact, if I am not wrong in the past the topk was visible in the documentation as a global pooling method.

About the sorting thing, I will try the topk with ratio=1.0 and see if it works.

[rusty1s]

Do you mean about making visible the topk?

Yes :)