选择最大值及其索引,大家都知道使用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。

Logo

为开发者提供学习成长、分享交流、生态实践、资源工具等服务,帮助开发者快速成长。

更多推荐