现在位置: 首页 > PyTorch 教程 > 正文

PyTorch torch.cumsum 函数


Pytorch torch 参考手册 Pytorch torch 参考手册

torch.cumsum 是 PyTorch 中用于计算累积和的函数。它返回沿指定维度的累积和,即从开始到当前位置所有元素的和。

函数定义

torch.cumsum(input, dim, dtype=None)

使用示例

实例

import torch

# 计算累积和
x = torch.tensor([1, 2, 3, 4, 5])
result = torch.cumsum(x, dim=0)
print("输入:", x)
print("累积和:", result)
# 输出: tensor([1, 3, 6, 10, 15])
# 说明: 1, 1+2=3, 1+2+3=6, 1+2+3+4=10, 1+2+3+4+5=15

# 2 维张量
x = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.float32)
# 按列累积
result_col = torch.cumsum(x, dim=0)
print("n按列累积和:")
print(result_col)
# tensor([[ 1,  2,  3],
#         [ 5,  7,  9]])

# 按行累积
result_row = torch.cumsum(x, dim=1)
print("n按行累积和:")
print(result_row)
# tensor([[ 1,  3,  6],
#         [ 4,  9, 15]])

Pytorch torch 参考手册 Pytorch torch 参考手册