torch.sum(input, dim=None, keepdim=False, dtype=None)
input
: 要求和的输入张量。dim
(可选): 要沿着哪个维度求和。如果不指定,将对整个张量进行求和,返回一个标量。keepdim
(可选): 如果设置为True
,结果张量将保持与输入张量相同的维度数,但求和维度的大小将变为1。默认值为False
。dtype
(可选): 返回的张量的数据类型。如果不指定,将使用输入张量的数据类型。
1. 对整个张量求和,返回标量:
import torch
x = torch.tensor([[1, 2, 3], [4, 5, 6]])
total_sum = torch.sum(x)
print(total_sum) # 输出: tensor(21)
2.沿着指定的维度求和,保持维度不变:
import torch
x = torch.tensor([[1, 2, 3], [4, 5, 6]])
sum_along_dim0 = torch.sum(x, dim=0)
print(sum_along_dim0) # 输出: tensor([5, 7, 9])
sum_along_dim1 = torch.sum(x, dim=1)
print(sum_along_dim1) # 输出: tensor([ 6, 15])
3.使用 keepdim=True
保持维度:
import torch
x = torch.tensor([[1, 2, 3], [4, 5, 6]])
sum_along_dim0_with_keepdim = torch.sum(x, dim=0, keepdim=True)
print(sum_along_dim0_with_keepdim) # 输出: tensor([[5, 7, 9]])
sum_along_dim1_with_keepdim = torch.sum(x, dim=1, keepdim=True)
print(sum_along_dim1_with_keepdim)
# 输出:
# tensor([[ 6],
# [15]])
4.指定数据类型:
import torch
x = torch.tensor([[1, 2, 3], [4, 5, 6]])
sum_along_dim0_with_keepdim = torch.sum(x, dim=0, keepdim=True)
print(sum_along_dim0_with_keepdim) # 输出: tensor([[5, 7, 9]])
sum_along_dim1_with_keepdim = torch.sum(x, dim=1, keepdim=True)
print(sum_along_dim1_with_keepdim)
# 输出:
# tensor([[ 6],
# [15]])