torch.topk()函数快速理解

函数作用：

该函数的作用即按字面意思理解，topk:取数组的前k个元素进行排序。

通常该函数返回2个值，第一个值为排序的数组，第二个值为该数组中获取到的元素在原数组中的位置标号。

---

举个栗子：

import numpy as np
import torch
import torch.utils.data.dataset as Dataset
from torch.utils.data import Dataset,DataLoader

####################准备一个数组#########################
tensor1=torch.tensor([[10,1,2,1,1,1,1,1,1,1,10],
             [3,4,5,1,1,1,1,1,1,1,1],
             [7,8,9,1,1,1,1,1,1,1,1],
             [1,4,7,1,1,1,1,1,1,1,1]],dtype=torch.float32)

####################打印这个原数组#########################
print('tensor1:')
print(tensor1)

#################使用torch.topk()这个函数##################
print('使用torch.topk()这个函数得到：')

'''k=3代表从原数组中取得3个元素，dim=1表示从原数组中的第一维获取元素
（在本例中是分别从[10,1,2,1,1,1,1,1,1,1,10]、[3,4,5,1,1,1,1,1,1,1,1]、
  [7,8,9,1,1,1,1,1,1,1,1]、[1,4,7,1,1,1,1,1,1,1,1]这四个数组中获取3个元素）
其中largest=True表示从大到小取元素'''
print(torch.topk(tensor1, k=3, dim=1, largest=True))


#################打印这个函数第一个返回值####################
print('函数第一个返回值topk[0]如下')
print(torch.topk(tensor1, k=3, dim=1, largest=True)[0])

#################打印这个函数第二个返回值####################
print('函数第二个返回值topk[1]如下')
print(torch.topk(tensor1, k=3, dim=1, largest=True)[1])
'''

#######################运行结果##########################
tensor1:
tensor([[10.,  1.,  2.,  1.,  1.,  1.,  1.,  1.,  1.,  1., 10.],
        [ 3.,  4.,  5.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.],
        [ 7.,  8.,  9.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.],
        [ 1.,  4.,  7.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.]])

使用torch.topk()这个函数得到：

'得到的values是原数组dim=1的四组从大到小的三个元素值；
得到的indices是获取到的元素值在原数组dim=1中的位置。'


torch.return_types.topk(
values=tensor([[10., 10.,  2.],
        [ 5.,  4.,  3.],
        [ 9.,  8.,  7.],
        [ 7.,  4.,  1.]]),
indices=tensor([[ 0, 10,  2],
        [ 2,  1,  0],
        [ 2,  1,  0],
        [ 2,  1,  0]]))

函数第一个返回值topk[0]如下
tensor([[10., 10.,  2.],
        [ 5.,  4.,  3.],
        [ 9.,  8.,  7.],
        [ 7.,  4.,  1.]])
        
函数第二个返回值topk[1]如下
tensor([[ 0, 10,  2],
        [ 2,  1,  0],
        [ 2,  1,  0],
        [ 2,  1,  0]])
'''

---

该函数功能经常用来获取张量或者数组中最大或者最小的元素以及索引位置，是一个经常用到的基本函数。