对pytorch中 torch.argmax（dim=）、x.argmin(dim=)的容易理解

import torch

x=torch.randn(3,3,4)
print(x)

print(x.argmax(dim=0))



输出：tensor([[[ 0.5128,  0.3717,  0.3606, -0.0286],
         [ 0.0933, -1.4781, -0.3561, -0.2652],
         [-0.8861,  0.6988,  1.1243, -1.1301]],

        [[-0.0246,  0.0917, -0.0623, -1.4874],
         [-2.0169, -0.8390,  0.1292, -0.1190],
         [-0.4090, -0.9277, -0.9193,  2.1597]],

        [[ 0.7030, -1.3322, -2.0856,  0.3122],
         [ 1.4656, -1.1145,  0.8208, -1.8250],
         [-0.8721,  0.4831,  1.5668,  0.1656]]])




tensor([[2, 0, 0, 2],
        [2, 1, 2, 1],
        [1, 0, 2, 1]])

（1）dim=0,比较的是torch.randn(3,3,4)中第一个第一个3所在维度的比较，也就是：

        [ 0.5128,  0.3717,  0.3606, -0.0286],
         [ 0.0933, -1.4781, -0.3561, -0.2652],                              下标为0
         [-0.8861,  0.6988,  1.1243, -1.1301]

        [-0.0246,  0.0917, -0.0623, -1.4874],
         [-2.0169, -0.8390,  0.1292, -0.1190],                             下标为1
         [-0.4090, -0.9277, -0.9193,  2.1597]]

        [ 0.7030, -1.3322, -2.0856,  0.3122],
         [ 1.4656, -1.1145,  0.8208, -1.8250],                           下标为2
         [-0.8721,  0.4831,  1.5668,  0.1656]]

三个矩阵（下标从0开始，）对应元素的比较，

所以print(x.argmax(dim=0))，就是0.5128和-0.0246和0.7030比较，输出最大值所在的下标（0.7030所在的张量下标为2，所以输出为2）；同理0.3717和0.0917和-1.3322进行比较，输出最大值的下标0（最大值为0.3717，下标为0）；以此类推。

所以dim=0 的输出为：

tensor([[2, 0, 0, 2],
            [2, 1, 2, 1],
           [1, 0, 2, 1]])   。所以看到输出你就很容易理解那个矩阵对应元素的值最大。

（2）dim=1,及（3,3,4）对第二个维度的元素进行比较（第二个3）。

import torch

x=[[[ 0.5128,  0.3717,  0.3606, -0.0286],
         [ 0.0933, -1.4781, -0.3561, -0.2652],
         [-0.8861,  0.6988,  1.1243, -1.1301]],

        [[-0.0246,  0.0917, -0.0623, -1.4874],
         [-2.0169, -0.8390,  0.1292, -0.1190],
         [-0.4090, -0.9277, -0.9193,  2.1597]],

        [[ 0.7030, -1.3322, -2.0856,  0.3122],
         [ 1.4656, -1.1145,  0.8208, -1.8250],
         [-0.8721,  0.4831,  1.5668,  0.1656]]]
x=torch.tensor(x)



print(x.argmax(dim=1))


输出：
tensor([[0, 2, 2, 0],
        [0, 0, 1, 2],
        [1, 2, 2, 0]])

        [ 0.5128,  0.3717,  0.3606, -0.0286],                  行标为0
         [ 0.0933, -1.4781, -0.3561, -0.2652],                   行标为1          
         [-0.8861,  0.6988,  1.1243, -1.1301]                 行标为2

对行元素进行比较 0.5128，0.0933，-0.8861进行比较，输出最大值所在的下标，最大值为0.5128，行标为0，所以输出为0；0.3717，-1.4781，0.6988进行比较，输出最大值得下标2（0.6988最大，行下标为2），以此类推第一个矩阵输出[0, 2, 2, 0]，第二个矩阵输出

 [0, 0, 1, 2]，第三个矩阵输出[1, 2, 2, 0]。即dim=1输出tensor([[0, 2, 2, 0],
                                                                                                     [0, 0, 1, 2],
                                                                                                     [1, 2, 2, 0]])

（3）dim=2，（3,3,4）对第三个维度的元素进行比较（4）。

import torch

x=[[[ 0.5128,  0.3717,  0.3606, -0.0286],
         [ 0.0933, -1.4781, -0.3561, -0.2652],
         [-0.8861,  0.6988,  1.1243, -1.1301]],

        [[-0.0246,  0.0917, -0.0623, -1.4874],
         [-2.0169, -0.8390,  0.1292, -0.1190],
         [-0.4090, -0.9277, -0.9193,  2.1597]],

        [[ 0.7030, -1.3322, -2.0856,  0.3122],
         [ 1.4656, -1.1145,  0.8208, -1.8250],
         [-0.8721,  0.4831,  1.5668,  0.1656]]]
x=torch.tensor(x)



print(x.argmax(dim=2))

输出：
tensor([[0, 0, 2],
        [1, 2, 3],
        [0, 0, 2]])

           第0列       第1列     第2列     第3列

         [ 0.5128,  0.3717,  0.3606, -0.0286],                 
         [ 0.0933, -1.4781, -0.3561, -0.2652],                            
         [-0.8861,  0.6988,  1.1243, -1.1301]       

  第一行 [ 0.5128,  0.3717,  0.3606, -0.0286]中最大的数0.5128所在的列标号为0，

第二行   [ 0.0933, -1.4781, -0.3561, -0.2652]最大元素0.0933所在的列标号为0，

第三行 [-0.8861,  0.6988,  1.1243, -1.1301]    最大元素1.1243所在的列标号为2，

所以输出[0, 0, 2]

        [-0.0246,  0.0917, -0.0623, -1.4874],
         [-2.0169, -0.8390,  0.1292, -0.1190],
         [-0.4090, -0.9277, -0.9193,  2.1597]

同理，输出[1, 2, 3]，

        [ 0.7030, -1.3322, -2.0856,  0.3122],
         [ 1.4656, -1.1145,  0.8208, -1.8250],
         [-0.8721,  0.4831,  1.5668,  0.1656]

输出[0, 0, 2]

所以dim=2，总的输出为tensor([[0, 0, 2],
                                                     [1, 2, 3],
                                                     [0, 0, 2]])