在原文章上有自己的补充和修改
【参考:torch.argmax函数说明_Egozjuer的博客-CSDN博客】
argmax函数:torch.argmax(input, dim=None, keepdim=False)返回指定维度最大值的序号
dim给定的定义是:the demention to reduce.
也就是把在dim这个维度的数据,变成这个维度上的最大值的下标index
1)dim的不同值表示不同维度。特别的在dim=0表示二维中的列,dim=1在二维矩阵中表示行。广泛的来说,我们不管一个矩阵是几维的,比如一个矩阵维度如下:(d0,d1,…,dn−1) ,那么dim=0就表示对应到d0 也就是第一个维度,dim=1表示对应到第二个维度,以此类推。
2)知道dim的值是什么意思还不行,还要知道函数中这个dim给出来会发生什么。
结合这两个才会知道dim在函数中的作用。下面举两个例子说明上面的第二点。
例子一:torch.argmax()函数中dim表示该维度会消失
这个消失是什么意思?官方英文解释是:dim (int) – the dimension to reduce.
我们知道argmax就是得到最大值的序号索引,
对于一个维度为(d0,d1) 的矩阵来说,我们想要求每一行中最大数的在该行中的列号,最后我们得到的就是一个维度为(d0,1) 的一矩阵。这时候,列就要消失了。
因此,我们想要求每一行最大的列标号,我们就要指定dim=1,表示我们不要列了,保留行的size就可以了。
假如我们想求每一列的最大行标,就可以指定dim=0,表示我们不要行了。
- 二维数据
import torch
# (3,4)
a=torch.tensor(
[
[1, 5, 5, 2],
[9, -6, 2, 8],
[-3, 7, -9, 1]
])
b=torch.argmax(a,dim=0)
print(b)
"""
tensor([1, 2, 0, 1]) # (1,4)
dim=0 跨行 => 每列
(多行 => 一行)
第一列:1 9 -3 : 9最大,下标为1
第二列:5 -6, 7 : 7最大,下标为2
以此类推
"""
理解:dim=0维度上为3,即在那3组数据中作比较,求得是每一列中的最大行标,因此为[1,2,0,4]
- 三维数据
import torch
# (2,3,4)
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)
"""
tensor([[0, 1, 0, 1],
[1, 1, 1, 1],
[1, 1, 1, 1]])
1 与 -1 比较,1大,1的下标为0
"""
dim=0,即将第一个维度消除,维度为2,也就是将两个[34]矩阵只保留一个,因此要在两组中作比较,即将上下两个[34]的矩阵分别在对应的位置上比较
b=torch.argmax(a,dim=1)
"""
tensor([[1, 2, 0, 1],
[1, 0, 2, 1]])
"""
dim=1,即将第二个维度消除,维度为3,在那3组数据中作比较,求得是每一列中的最大行标
这么理解:矩阵维度变为[2*4];
b=torch.argmax(a,dim=2)
"""
tensor([[1, 0, 1],
[1, 0, 2]])
"""
dim=2,即将第三个维度消除,维度为4,在那4组数据中作比较,求得是每一行中的最大列标
这么理解:矩阵维度变为[2*3]