torch.sum的用法

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