pytorch选出数据中的前k个最大(最小)值及其索引
选择最大值及其索引,大家都知道使用max(),argmax()函数。那么如何返回前k个最大值呢,这在我们计算topK准确率的时候很有必要:在torch中,我们可以使用sort函数来实现:a, idx1 = torch.sort(data, descending=True)idx = idx1[:k]Return:a:排好序的数据idx1:对应排序数据的索引因此只需设置k的大小,就可以截取到前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。
更多推荐
已为社区贡献3条内容
所有评论(0)