导读
最近有时间看一些目标检测项目的代码（基于Pytorch），里边很多Pytorch的相关操作都忘记了，特来此记录一下，用以加深记忆，而且还能以备一样处境的同学前来查询。今天的主角是torch.argmax(input, dim, keepdim=False)。
官方文档地址
https://pytorch.org/docs/stable/generated/torch.argmax.html
torch.argmax(input) → LongTensor
Returns the indices of the maximum value of all elements in the input tensor.
根据官方的解释，该函数可以返回输入张量中所有元素的最大值的索引。当然这只是最初级的用法，根据输入参数的不同，其返回的结果也不同。下面我们一起了解它的参数都有哪些作用。
参数解析
Parameters
- input (Tensor) ： the input tensor.
- dim (int) ：the dimension to reduce. If None, the argmax of the flattened input is returned.
- keepdim (bool) ： whether the output tensor has dim retained or not. Ignored if dim=None.
这是官网上对参数的解释，input就是我们输入的要操作的张量；dim是我们选择的要在张量的哪个维度上进行计算，输出这个维度最大值的索引，这里一行元素变成一个索引，所以官网中用了reduce；keepdim是询问输出是否与输入保持一样的形状，默认是不保持（False）。
举例演示
首先输入一个张量，注意我们输入的这个张量的shape为[5, 9]
>>> x = torch.randn(5,9)
>>> print(x)
tensor([[ 0.3918, 0.3978, 0.2819, -0.8487, -1.0499, 0.2124, -1.3527, -1.5335,
1.1050],
[ 0.8450, -0.3717, -0.4705, -0.4024, 2.1019, -0.8545, 1.9085, 0.5792,
-0.4279],
[ 0.1993, -0.2887, 0.4467, 0.4878, 1.4934, -1.3862, 0.3576, -0.2363,
-2.0700],
[ 0.0536, 0.9385, 1.2661, -0.3469, -0.5772, -0.7822, 0.8315, -1.7256,
-0.4979],
[ 1.1592, -0.1604, 0.2798, 0.5974, 0.1782, -2.3354, -1.7775, -0.8366,
1.8993]])
接下来，现在第一个维度上进行操作
>>> torch.argmax(x,dim=0)
tensor([4, 3, 3, 4, 1, 0, 1, 1, 4])
第一个维度是行，即按行计算，我们看到结果输出的维度为9，正好是输入张量x的列数。torch.argmax()的计算方式如下：
每次在所有行的相同位置取元素，然后计算取得元素集合的最大值索引。
第一次取所有行的第一位元素，x[:, 0], 得到
tensor([0.3918, 0.8450, 0.1993, 0.0536, 1.1592])
第二次取所有行的第二位元素，x[:, 1], 得到
tensor([0.3978, -0.3717, -0.2887, 0.9385, -0.1604])
依次类推，x有9列，我们也可以取9次，所有取的结果如下：
tensor([ 0.3918, 0.8450, 0.1993, 0.0536, 1.1592])
tensor([ 0.3978, -0.3717, -0.2887, 0.9385, -0.1604])
tensor([ 0.2819, -0.4705, 0.4467, 1.2661, 0.2798])
tensor([-0.8487, -0.4024, 0.4878, -0.3469, 0.5974])
tensor([-1.0499, 2.1019, 1.4934, -0.5772, 0.1782])
tensor([-1.3527, 1.9085, 0.3576, 0.8315, -1.7775])
tensor([-1.5335, 0.5792, -0.2363, -1.7256, -0.8366])
tensor([ 1.1050, -0.4279, -2.0700, -0.4979, 1.8993])
然后分别计算以上每个张量中元素的最大值的索引，便得到tensor([4, 3, 3, 4, 1, 0, 1, 1, 4])
同理，按照列来操作也是一样的思路，这里就不详细说了，看结果：
>>> torch.argmax(x,dim=1)
tensor([8, 4, 4, 2, 8])
经过上边例子的演示，我们可以知道torch.argmax(input，dim)可以返回input中dim维度上的最大值索引。
我们给x在目标检测中赋予具体的含义，假如x的形状为[num_bbox, anchor]，那么x便是5个预测框分别与9个anchor计算得到的交并比，我们要选出来与预测框交并比最大的那个anchor，用来回归预测框越来越接近GT。这时候就要用到torch.argmax()找到与bbox交并比最大的anchor的序号。
>>> torch.argmax(x,dim=1)
tensor([8, 4, 4, 2, 8])
即与第一个预测框交并比最大的是第9个anchor，与第二个预测框交并比最大的是第5个anchor…