PyTorch torch.sum()求和函数的用法（附带实例）

PyTorch与其他统计软件一样，也内置了丰富的统计函数。使用 torch.sum() 函数求和，该函数对输入的张量数据的某一维度求和，一共有两种格式。

第一种格式：

torch.sum(input, dtype=None)

计算输入张量 input 中所有元素的和，返回一个标量值。

第二种格式：

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

计算输入张量 input 在指定维度 dim 上的和，并返回一个新的张量。

参数说明：

• input：输入一个张量。
• dim：要求和的维度，可以是一个列表。当 dim=0 时，即第 0 个维度会缩减，也就是说将N行压缩成一行，故相当于对列进行求和；当 dim=1 时，对行进行求和。
• keepdim：求和之后这个 dim 的元素个数为 1，所以要被去掉，如果要保留这个维度，则要保证 keepdim=True。如果 keepdim 为 True，则保留原始张量的维度；如果 keepdim 为 False，则不保留原始张量的维度。

【实例】首先，创建初始张量，代码如下：

# 导入torch库
import torch
# 创建一个2×2的张量a，元素为随机数
a = torch.rand(2,2)
print(a)

输出如下：

tensor([[0.0528, 0.3420],
       [0.5011, 0.4264]])

设置参数 input 和 dim，代码如下：

a1 = torch.sum(a)                # 计算张量a所有元素的和
a2 = torch.sum(a, dim=(0, 1))    # 沿着第0维和第1维计算张量a的元素和
a3 = torch.sum(a, dim=0)         # 沿着第0维计算张量a的元素和
a4 = torch.sum(a, dim=1)         # 沿着第1维计算张量a的元素和
print(a1)
print(a2)
print(a3)
print(a4)

输出如下：

tensor(1.3223)
tensor(1.3223)
tensor([0.5539, 0.7684])
tensor([0.3948, 0.9275])

设置参数 keepdim，代码如下：

a5 = torch.sum(a, dim=(0, 1), keepdim=True) # 沿着第0维和第1维计算张量a的元素和，保持原始维度
a6 = torch.sum(a, dim=(0, ), keepdim=True)  # 沿着第0维计算张量a的元素和，保持原始维度
a7 = torch.sum(a, dim=(1, ), keepdim=True)  # 沿着第1维计算张量a的元素和，保持原始维度
print(a5)
print(a6)
print(a7)

输出如下：

tensor([[1.3223]])
tensor([[0.5539, 0.7684]])
tensor([[0.3948],
[0.9275]])