torch.sum()函数用法

在 PyTorch 中，torch.sum() 用于计算输入张量的所有元素之和。它的语法如下：

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

其中，参数含义如下：

• input：要进行处理的张量。
• dim：指定沿着哪个维度求和。如果为 None，则对整个张量求和。默认值为 None。
• keepdim：当设为 True 时，输出张量形状会保持与输入张量相同；否则输出张量会自动压缩删除所有大小为 1 的维度。默认为 False。
• dtype：返回张量的数据类型。可选值有 float32、float64、int32 等。

例如，给定一个张量 a，可以使用 torch.sum() 函数计算出张量 a 中所有元素的和：

import torch

a = torch.tensor([[1, 2], [3, 4]])
b = torch.sum(a)

print(b)  # 输出：tensor(10)

python复制代码

import torch a = torch.tensor([[1, 2], [3, 4]]) b = torch.sum(a) print(b) # 输出：tensor(10)

在上面的例子中，我们首先定义了一个形状为 (2, 2) 的张量，并将其存储到变量 a 中，然后使用 torch.sum() 函数计算出 a 中所有元素的和，最终得到 tensor(10)。

此外，我们还可以通过设置 dim 的值来求解沿着张量的某一个维度的和。例如：

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

a = torch.tensor([[1, 2], [3, 4]])
b = torch.sum(a)

print(b)  # 输出：tensor(10)

a = torch.tensor([[1, 2], [3, 4]])
b = torch.sum(a, dim=1)

print(b)  # 输出：tensor([3, 7])

在上面的例子中，我们通过设置 dim=1 的方式来计算沿着变量 a 的第二个维度 (即每一行）的和，并返回一个形状为 (2,) 的张量。