非常棒的直观方法_理解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