pytorch中topk()用法的测试与个人理解

参数介绍：

直接官网的介绍topk()

• input：就是输入的tensor，也就是要取topk的张量
• k：就是取前k个最大的值。
• dim：就是在哪一维来取这k个值。
• lagest：默认是true表示取前k大的值，false则表示取前k小的值
• sorted：是否按照顺序输出，默认是true。
• out ： 可选输出张量 (Tensor, LongTensor)

直接上代码：

---

首先研究一下dim和k这两个最重要的参数：

import torch
seed = 0
torch.manual_seed(seed)
a = torch.randint(1,10,(3,4,4))

print(a)

tensor([[[9, 1, 3, 7],
         [8, 7, 8, 2],
         [2, 1, 9, 3],
         [7, 4, 2, 3]],

        [[1, 1, 6, 4],
         [9, 3, 9, 3],
         [9, 6, 8, 9],
         [7, 1, 2, 1]],

        [[9, 7, 2, 7],
         [2, 9, 9, 8],
         [3, 4, 8, 8],
         [2, 6, 5, 8]]])

values , indices = a.topk(2,dim=0)
print(values.shape)
print(values)
print(indices.shape)
print(indices)

torch.Size([2, 4, 4])
tensor([[[9, 7, 6, 7],
         [9, 9, 9, 8],
         [9, 6, 9, 9],
         [7, 6, 5, 8]],

        [[9, 1, 3, 7],
         [8, 7, 9, 3],
         [3, 4, 8, 8],
         [7, 4, 2, 3]]])
torch.Size([2, 4, 4])
tensor([[[0, 2, 1, 0],
         [1, 2, 1, 2],
         [1, 1, 0, 1],
         [0, 2, 2, 2]],

        [[2, 0, 0, 2],
         [0, 0, 2, 1],
         [2, 2, 1, 2],
         [1, 0, 0, 0]]])

**这里说一下我的理解 **

首先dim=0这个参数表示在某一维取topk，在我的代码中就是取前2个。首先看输出的values和indices的张量形状：（2，4，4）这里可以结合下面两个实验总结出规律，dim取几，输出结果的形状就是：其他维度不变，对应维度变成k。

现在dim=0最后的输出就是要变成（2，4，4）也就是之前第一维中保留两个最大的。

看values的值，讲一下第一行元素[9,7,6,7]是如何得来的：
因为dim=0所以要从第0维来看，将数据分成3份，分别是：

1. [[9, 1, 3, 7],
    [8, 7, 8, 2],
    [2, 1, 9, 3],
    [7, 4, 2, 3]]
2. [[1, 1, 6, 4],
    [9, 3, 9, 3],
    [9, 6, 8, 9],
    [7, 1, 2, 1]]
3. [[9, 7, 2, 7],
    [2, 9, 9, 8],
    [3, 4, 8, 8],
    [2, 6, 5, 8]]

要以这三个tensor为单位进行topk的筛选，首先比较每一个tensor的第一行，因为参数k为2，所以就要找到这3组元素中的最大值和次大值，作为最后的输出。因此最大值就是[9，7，6，7]，次大值为:[9，1，3，7]这样就完成了筛选。索引值也就是当前位置处的元素，是来自于这三个元素中的哪一个。我认为把这个看懂后面就可以迎刃而解，大家可以仔细理解一下不太懂的话也没关系，看完后面两个可能这个就懂了。

values1 , indices1 = a.topk(2,dim=1)
print(values1)
print(indices1)

torch.Size([3, 2, 4])
tensor([[[9, 7, 9, 7],
         [8, 4, 8, 3]],

        [[9, 6, 9, 9],
         [9, 3, 8, 4]],

        [[9, 9, 9, 8],
         [3, 7, 8, 8]]])
torch.Size([3, 2, 4])
tensor([[[0, 1, 2, 0],
         [1, 3, 1, 2]],

        [[1, 2, 1, 2],
         [2, 1, 2, 0]],

        [[0, 1, 1, 2],
         [2, 0, 2, 3]]])

这个例子是dim=1时，类比于dim=0的情况。这里是对第一维进行筛选操作。需要注意的是这里第0维的三个元素是分开操作的。这里我提供一种我自己的理解思路大家借鉴。首先还是按照第0维将tensor分为3块

1. [[9, 1, 3, 7],
    [8, 7, 8, 2],
    [2, 1, 9, 3],
    [7, 4, 2, 3]]
2. [[1, 1, 6, 4],
    [9, 3, 9, 3],
    [9, 6, 8, 9],
    [7, 1, 2, 1]]
3. [[9, 7, 2, 7],
    [2, 9, 9, 8],
    [3, 4, 8, 8],
    [2, 6, 5, 8]]

这里每一块中的第0维就是总体tensor的第一维，从第0维来看就是4个14的向量，因此就是对这4向量取最大值和次大值。也就是在这个44的张量中选出对应位置的最大值和次大值。例如第一块中筛选出的结果就是[9，7，9，7]和[8，4，8，3]其他同理，索引值表示当前位置处的值是来自哪一个向量。

values2 , indices2 = a.topk(2,dim=2)
print(values2.shape)
print(values2)
print(indices2.shape)
print(indices2)

torch.Size([3, 4, 2])
tensor([[[9, 7],
         [8, 8],
         [9, 3],
         [7, 4]],

        [[6, 4],
         [9, 9],
         [9, 9],
         [7, 2]],

        [[9, 7],
         [9, 9],
         [8, 8],
         [8, 6]]])
torch.Size([3, 4, 2])
tensor([[[0, 3],
         [0, 2],
         [2, 3],
         [0, 1]],

        [[2, 3],
         [0, 2],
         [3, 0],
         [0, 2]],

        [[0, 1],
         [1, 2],
         [2, 3],
         [3, 1]]])

类比前两种情况的思考方式，这里的操作就是对整个张量最内层做的操作，也就是整体张量形状(3,4,4)中的4这个4就是最内层每一个一维向量中的4个元素，取对应的最大值和次大值，应该也容易理解。大家可以对比着三种情况的输入输出加以理解。
另外，k参数默认是最后一维

然后研究一下lagest参数：
直接用最后一维

values2 , indices2 = a.topk(2,dim=2,largest=False)
print(values2.shape)
print(values2)
print(indices2.shape)
print(indices2)

torch.Size([3, 4, 2])
tensor([[[1, 3],
         [2, 7],
         [1, 2],
         [2, 3]],

        [[1, 1],
         [3, 3],
         [6, 8],
         [1, 1]],

        [[2, 7],
         [2, 8],
         [3, 4],
         [2, 5]]])
torch.Size([3, 4, 2])
tensor([[[1, 2],
         [3, 1],
         [1, 0],
         [2, 3]],

        [[1, 0],
         [1, 3],
         [1, 2],
         [1, 3]],

        [[2, 3],
         [0, 3],
         [0, 1],
         [0, 2]]])

很明显，只是取最小和次小

---

下面是sorted参数
依然用dim=2进行测试

values2 , indices2 = a.topk(2,dim=2,sorted=False)
print(values2.shape)
print(values2)
print(indices2.shape)
print(indices2)

torch.Size([3, 4, 2])
tensor([[[9, 7],
         [8, 8],
         [9, 3],
         [7, 4]],

        [[6, 4],
         [9, 9],
         [9, 9],
         [7, 2]],

        [[9, 7],
         [9, 9],
         [8, 8],
         [8, 6]]])
torch.Size([3, 4, 2])
tensor([[[0, 3],
         [0, 2],
         [2, 3],
         [0, 1]],

        [[2, 3],
         [0, 2],
         [3, 0],
         [0, 2]],

        [[0, 1],
         [1, 2],
         [2, 3],
         [3, 1]]])

这里和不sorted=True好像并没有区别，不知道要怎么理解，网上也没找到类似的解释，希望有知道的大佬可以多多指教！！

如有错误请多指正