【Pytorch】torch.argmax()用法

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

import numpy as np
import torch
b = torch.tensor([
    [
        [8, 10, 14, 21],
        [9, 6, 23, 13],
        [15, 32,16,11]
    ],

    [
        [31,20, 27, 17],
        [28, 34, 22,33],
        [36, 30, 1, 3]
    ],

    [
        [12, 29, 26, 25],
        [19, 7, 5, 4],
        [2, 35, 24, 18]
    ]
])
print('b.shape=',b.shape) #b.shape=([3,3,4])
dim_0 = torch.argmax(b,dim = 0)
print('dim_0.shape=',dim_0.shape)
print('dim_0=',dim_0)

dim_1 = torch.argmax(b,dim = 1)
print('dim_1.shape=',dim_1.shape)
print('dim_1=',dim_1)

dim_2 = torch.argmax(b,dim = 2)
print('dim_2.shape=',dim_2.shape)
print('dim_2=',dim_2)

dim_f1 = torch.argmax(b,dim = -1)
print('dim_-1.shape=',dim_f1.shape)
print('dim_-1=',dim_f1)

dim_f2 = torch.argmax(b,dim = -2)
print('dim_-2.shape=',dim_f2.shape)
print('dim_-2=',dim_f2)

代码运行结果：

b.shape= torch.Size([3, 3, 4])
dim_0.shape= torch.Size([3, 4])
dim_0= tensor([[1, 2, 1, 2],
        [1, 1, 0, 1],
        [1, 2, 2, 2]])
dim_1.shape= torch.Size([3, 4])
dim_1= tensor([[2, 2, 1, 0],
        [2, 1, 0, 1],
        [1, 2, 0, 0]])
dim_2.shape= torch.Size([3, 3])
dim_2= tensor([[3, 2, 1],
        [0, 1, 0],
        [1, 0, 1]])
dim_-1.shape= torch.Size([3, 3])
dim_-1= tensor([[3, 2, 1],
        [0, 1, 0],
        [1, 0, 1]])
dim_-2.shape= torch.Size([3, 4])
dim_-2= tensor([[2, 2, 1, 0],
        [2, 1, 0, 1],
        [1, 2, 0, 0]])

进程已结束，退出代码为 0

总结：
dim=0,将张量最高维度消除，也就是说将b张量为([3 ,3,4])变为([3,4]);
同理，dim=1，将第二高的维度消除，也就是说将b张量为([3,3 ,4])变为([3,4])；dim=2，将第三高的维度消除，也就是说将b张量为([3,3,4 ])变为([3,3])；以此类推。dim=-1表示张量维度的最低维度 -2表示张量的倒数第二维度，-3表示倒数第三维度。
下一篇：
【Pytorch】F.softmax()方法说明