关于函数torch.topk用法的思考

开始介绍之前先来点哲理性的思考，为什么函数 torch.topk，他的名字会叫 topk呢？

个人认为名称是来源于 “top k”，在这种情况下，它表示 “前 k 个最大值”。

假设我们有一个形状为 $(2,3,4)$ 的三维张量 $A$，如下所示：

A = torch.tensor([[[ 1,  3,  5,  7],
                   [ 2,  4,  6,  8],
                   [ 9, 11, 13, 15]],
                  [[16, 18, 20, 22],
                   [17, 19, 21, 23],
                   [10, 12, 14, 24]]])

1. 沿着dim=0

沿着 dim=0（即在子矩阵之间进行比较）：

k = 1
topk_values, topk_indices = torch.topk(A, k=k, dim=0)

A = torch.tensor([[[ 1,  3,  5,  7],
                   [ 2,  4,  6,  8],
                   [ 9, 11, 13, 15]],
                  [[16, 18, 20, 22],
                   [17, 19, 21, 23],
                   [10, 12, 14, 24]]])

沿着 dim=0（即在子矩阵之间进行比较）：

topk_values = tensor([[[16, 18, 20, 22],
                        [17, 19, 21, 23],
                        [10, 12, 14, 24]]])

topk_indices = tensor([[[1, 1, 1, 1],
                         [1, 1, 1, 1],
                         [1, 1, 1, 1]]])

那此时我们令k = 3会发生什么？很显然，我们并没有三个子矩阵，所以此时程序会报错。给大家看一下程序的错误：

Traceback (most recent call last):
  File "E:\Learning_Material\Junior_Second_Semester\Adademic_Research\BusterNet_pytorch-master\test_2023_4_7.py", line 14, in <module>
    topk_values, topk_indices = torch.topk(A, k=k, dim=0)
RuntimeError: selected index k out of range

2. 沿着dim=1

沿着 dim=1（即在行之间进行比较）：

k = 2
topk_values, topk_indices = torch.topk(A, k=k, dim=1)

topk_values = tensor([[[ 9, 11, 13, 15],
         			   [ 2,  4,  6,  8]],

        			   [[17, 19, 21, 24],
         				[16, 18, 20, 23]]])

topk_indices = tensor([[[2, 2, 2, 2],
         				[1, 1, 1, 1]],

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

在沿着行进行比较的情况下，比较是不会在子矩阵，也就是更高维度上发生的。仅仅在子矩阵的维度上比较一个子矩阵中最大的行数。

3. 沿着dim=2

沿着 dim=2（即在列之间进行比较）：

k = 2
topk_values, topk_indices = torch.topk(A, k=k, dim=2)

topk_values = tensor([[[ 7,  5],
         			   [ 8,  6],
         			   [15, 13]],

        			  [[22, 20],
         		   	   [23, 21],
         			   [24, 14]]])

topk_indices = tensor([[[3, 2],
         				[3, 2],
         				[3, 2]],

        			   [[3, 2],
         				[3, 2],
         				[3, 2]]])

沿着dim=2比较，此时就不会牵扯到更高的两个维度，只会在最后一个维度之内进行排序比较

4. 总结

没有什么别的经验，希望对大家有用就好！