pytorch选出数据中的前k个最大（最小）值及其索引

选择最大值及其索引，大家都知道使用max(),argmax()函数。

那么如何返回前k个最大值呢，这在我们计算topK准确率的时候很有必要：

在torch中，我们可以使用sort函数来实现：

a, idx1 = torch.sort(data, descending=True)#descending为alse，升序，为True，降序
idx = idx1[:k]

Return:
a:排好序的数据
idx1：对应排序数据的索引

因此只需设置k的大小，就可以截取到前k个最大值的索引。这里若数据是tensor则用torch，若是list或ndarray，可以用numpy。