pytorch

torch.cumsum和torch.cumprod的用法

import torch

a = torch.arange(0,6).view(2,3)
b = torch.cumsum(a,dim=0)
c = torch.cumsum(a,dim=1)
print('a=',a)
print('b=',b)
print('c=',c)

dim=0时，a的第一行保持不变，从第二行开始，第二行的每个元素都是前面行的元素的累加和
dim=1时，a的第一列保持不变，从第二列开始，第二列的每个元素都是前面列的元素的累加和

import torch

a = torch.arange(0,6).view(2,3)
b = torch.cumprod(a,dim=0)
c = torch.cumprod(a,dim=1)
print('a=',a)
print('b=',b)
print('c=',c)

dim=0时，a的第一行保持不变，从第二行开始，第二行的每个元素都是前面行的元素的累积
dim=1时，a的第一列保持不变，从第二列开始，第二列的每个元素都是前面列的元素的累积