選擇最大值及其索引，大家都知道使用max(),argmax()函數(shù)。

那么如何返回前k個(gè)最大值呢，這在我們計(jì)算topK準(zhǔn)確率的時(shí)候很有必要：

在torch中，我們可以使用sort函數(shù)來實(shí)現(xiàn)：

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

Return:
a:排好序的數(shù)據(jù)
idx1：對(duì)應(yīng)排序數(shù)據(jù)的索引

因此只需設(shè)置k的大小，就可以截取到前k個(gè)最大值的索引。這里若數(shù)據(jù)是tensor則用torch，若是list或ndarray，可以用numpy。

總結(jié)

以上是生活随笔為你收集整理的pytorch选出数据中的前k个最大（最小）值及其索引的全部?jī)?nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。