How to get indices of top-K values from a numpy array

Question

Let suppose I have probabilities from a Pytorch or Keras predictions and result is with the softmax function

from scipy.special import softmax
probs = softmax(np.random.randn(20,10),1) # 20 instances and 10 class probabilities
probs

I want to find top-5 indices from this numpy array. All I want to do is to run a loop on the results something like:

for index in top_5_indices:
    if index in result:
        print('Found')

I'll get if my results are in top-5 results.

Pytorch has top-k function and I have seen numpy.argpartition but I have no idea how to get this done?

Answer 1

A little more expensive, but argsort would do:

idx = np.argsort(probs, axis=1)[:,-5:]

If we are talking about pytorch:

probs = torch.from_numpy(softmax(np.random.randn(20,10),1))

values, idx = torch.topk(probs, k=5, axis=-1)

Answer 2

np.argpartition(probs,-5)[:,-5:]