torch.sum(),dim=0,dim=1解析

直接上代码：

# 定义一个2维张量a
a = torch.tensor([[1,2,3],[4,5,6]])
print(a)
a.shape

output:

tensor([[ 1,  2,  3],
        [ 4,  5,  6]])
torch.Size([2, 3])
# 2行3列，没毛病

dim=0，降维（纵向压缩）

b = torch.sum(a,dim=0)
print(b)
b.shape

output:

tensor([ 5,  7,  9])
torch.Size([3])

5 = 1 + 4
7 = 2 + 5
9 = 3 + 6

此时b为1维张量，是一个向量，且剩余元素个数与a的列数相等。

假想a是一个站着的弹簧，从上向下压缩它，就是sum(dim=0)的效果。

dim=1，降维（横向压缩）

c = torch.sum(a,dim=1)
print(c)
c.shape

output:

tensor([  6,  15])
torch.Size([2])

6 = 1 + 2 + 3
15 = 4 + 5 +6

此时c也为1维张量，是一个向量，且剩余元素个数与a的行数相等。

假想a是一个躺着的弹簧，从右向左压缩它，就是sum(dim=1)的效果。