非常棒的直观方法_理解Pytorch中的dim
首先,我们创建一个张量
x = torch.tensor([
[1,2,3],
[4,5,6]
])
x.shape
>>torch.Size([2, 3])
此时,创建了一个2行3列的张量,2即为dim=0,3即为dim=1,那么按照dim=0求和时,就是对dim=0这个维度进行挤压成一行,如下所示:
那么扩展到3维的情况如下:
x = torch.tensor([
[
[1,2,3],
[4,5,6]
],
[
[1,2,3],
[4,5,6]
],
[
[1,2,3],
[4,5,6]
]
])
此时,x的shape为torch.Size([3,2,3]),其中dim=0是3,dim=1是2(行),dim=2是3(列),那么此时,torch.sum(x,dim=0)为:
现在对行dim=1求和,结果时3个为[5,7,9]的张量:
对列dim=2求和,结果即为把每个向量都按照列进行挤压:
Understanding dimensions in PyTorch