torch.argmax中dim详解

torch.argmax()函数

argmax函数：torch.argmax(input, dim=None, keepdim=False)返回指定维度最大值的序号，dim给定的定义是：the demention to reduce.也就是把dim这个维度的，变成这个维度的最大值的index。

例如tensor(2, 3, 4)

dim=0，将第一维度去掉，这样结果为tensor(3, 4)。

import torch
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)
print(b)
print(a.shape)
"""
tensor([[0, 1, 0, 1],
        [1, 1, 1, 1],
        [1, 1, 1, 1]])
torch.Size([2, 3, 4])
"""
# 去掉第一维度，这样剩下两个3x4数组，将两个数组的对应位置进行比较，例如：a[0][0][0]和a[1][0][0]比较
#因为a[1][0][0]大，所以b[0][0][0]就是1,以此类推  这里的0和1表示对应的位置，第0个数组还是第1数组大

dim=1， 将第二维度去掉，取每一列的最大值。结果展示为tensor(2, 4)

import torch
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=1)
print(b)
print(a.shape)
"""
tensor([[1, 2, 0, 1],
        [1, 2, 2, 1]])
torch.Size([2, 3, 4])
"""
# 去掉第二维度，结果为是一个2x4，将每一个3x4数组，变成1x4数组，经过变化后a[0] = tensor([9, 7, 5, 8])
#取每一列的最大值，a[0]中第一列的最大值的行标为1， 第二列的最大值的行标为2，第三列的最大值行标为0，第4列的最大值行标为1
#所以最后输出[1, 2, 0, 1]
#以此类推，

dim=2, 将第三维度去掉，取每一行的最大值。结果展示为tensor(2, 3)

import torch
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=2)
print(b)
print(a.shape)
"""
tensor([[2, 0, 1],
        [1, 0, 2]])
torch.Size([2, 3, 4])
"""
# 去掉第三维度，结果为是一个2x3，将每一个3x4数组，变成3x1数组，就好像经过变化后a[0] = tensor([5, 9, 7]的转置)
#取每一行的最大值，a[0]中第一行的最大值的列标为2， 第二行的最大值的列标为0，第三行的最大值列标为1，
#所以最后输出[2, 0, 1]
#以此类推，