torch.argmax（)

torch.argmax(input, dim=None, keepdim=False)

转载至 ： https://blog.csdn.net/qq_46092061/article/details/120612102

torch.argmax(input, dim=None, keepdim=False)

• 返回指定维度最大值的序号index

• 函数中dim表示该维度会消失。dim=0，表示我们不要行了，求出每一列的最大值的下标;dim=1，表示我们不要列了，保留每一行的最大值的下标。

二维例子：a = torch.tensor(
          [
              [1, 5, 5, 2],
              [9, -6, 2, 8],
              [-3, 7, -9, 1]
          ])
   b = torch.argmax(a, dim=0)
   b = tensor([1, 2, 0, 1])

三维例子：
a = torch.tensor([
          [
              [1, 5, 5, 2],
              [9, -6, 2, 8],
              [-3, 7, -9, 1]
          ],

          [
              [-1, 7, -5, 2],
              [9, 6, 2, 8],
              [3, 7, 9, 1]
          ]])
          
b = torch.argmax(a, dim=0)
b = tensor([[0, 1, 0, 1],
        [1, 1, 1, 1],
        [1, 1, 1, 1]])
a.shape():torch.Size([2, 3, 4])
dim=0,即将第一个维度消除，也就是将两个[3*4]矩阵只保留一个，
因此要在两组中作比较，即将上下两个[3*4]的矩阵分别在对应的位置上比较大小，保存较大的下标

b = torch.argmax(a, dim=1)
#dim=1，即将第二个维度消除,这么理解：矩阵维度变为[2*4];
tensor([[1, 2, 0, 1],
     [1, 2, 2, 1]])
torch.Size([2, 3, 4])