PyTorch中torch.sum()函数的用法总结

torch.sum()对输入的tensor数据的某一维度求和，一共两种用法。
方法一：

torch.sum(input, *, dtype=None) → Tensor

• input：输入的张量

案例：

x = torch.randn(2, 3)
print(x)
y = torch.sum(x)
print(y)

输出结果：

tensor([[-0.2328,  1.4580,  0.7448],
        [-0.7813,  0.3045, -1.9038]])
tensor(-0.4107)
# -0.2328+1.4580+0.7448-0.7813+0.3045-1.9038 = -0.41059999999999963

方法二：

torch.sum(input, dim, keepdim=False, *, dtype=None) → Tensor

• input:输入的张量
• dim:求和的维度，可以是一个列表，也就是可以同时接收多个维度，并可同时在这些维度上进行指定操作。
• keepdim:默认为False，若keepdim=True，则返回的Tensor除dim之外的维度与input相同。因为求和之后这个dim的元素个数为1，所以要被去掉，如果要保留这个维度，则应当keepdim=True。
  案例：

x = torch.arange(0, 12).view(3, 4)
print(x)
y1 = torch.sum(x, dim=1)
print(y1)
y2 = torch.sum(x, dim=0)
print(y2)
y3 = torch.sum(x, dim=0, keepdim=True)
print(y3)

输出结果：

# x
tensor([[ 0,  1,  2,  3],
        [ 4,  5,  6,  7],
        [ 8,  9, 10, 11]])
# y1
tensor([ 6, 22, 38])
# y2
tensor([12, 15, 18, 21])
# y3
tensor([[12, 15, 18, 21]])