torch.sum()对输入的tensor数据的某一维度求和,一共两种用法。
方法一:
torch.sum(input, *, dtype=None) → Tensor
案例:
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
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]])