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()方法说明