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]])

你可能感兴趣的:(人工智能,机器学习,pytorch,python,深度学习)