pytorch中有很多自带函数,掌握好这些函数,程序写起来当然非常有效率。
在这些pytorch函数中,很多都有dim这个控制参数,但是我们很难明白这个含义是什么。
本文试着总结一下:
1)dim的不同值表示不同维度。特别的在dim=0表示二维中的行,dim=1在二维矩阵中表示列。广泛的来说,我们不管一个矩阵是几维的,比如一个矩阵维度如下:
(
d
0
,
d
1
,
.
.
.
,
d
n
−
1
)
(d_0, d_1, ..., d_{n-1})
(d0,d1,...,dn−1),那么dim=0就表示对应到
d
0
d_0
d0也就是第一个维度,dim=1表示对应到
d
1
d_1
d1也就是第二个维度,依次类推。
2)知道dim的值是什么意思还不行,还要知道函数中这个dim给出来会发生什么。
结合这两个才会知道dim在函数中的作用。下面举两个例子说明上面的第二点。
例子一:torch.argmax()函数中dim表示该维度会消失。
这个消失是什么意思?官方英文解释是:dim (int) – the dimension to reduce.
我们知道argmax就是得到最大值的序号索引,对于一个维度为
(
d
0
,
d
1
)
(d_0, d_1)
(d0,d1)的矩阵来说,我们想要求每一行中最大数的在该行中的列号,最后我们得到的就是一个维度为
(
d
0
,
1
)
(d_0, 1)
(d0,1)的一矩阵。这时候,列就要消失了。
因此,我们想要求每一行最大的列标号,我们就要指定dim=1,表示我们不要列了,保留行的size就可以了。
假如我们想求每一列的最大行标,就可以指定dim=0,表示我们不要行了。
import torch
import os
import numpy as np
# os.environ['CUDA_VISIBLE_DEVICES'] = '1'
a = torch.rand((3,4))
print(a.size())
print(a)
b = torch.argmax(a, dim=1)
print(b)
print(b.size())
输出:
torch.Size([3, 4])
tensor([[0.8338, 0.6953, 0.7558, 0.5803],
[0.2105, 0.7638, 0.0912, 0.3341],
[0.5585, 0.8019, 0.6590, 0.2268]])
tensor([0, 1, 1])
torch.Size([3])
可以看见指定dim=1时,列的size没有了。