torch.cumprod实现累乘计算

木盏2023-11-08 0:00

cumprod取自"cumulative product"的缩写，即"累计乘法 "。

数学公式为：
y i = x 1 × x 2 × x 3 × . . . × x i y_i=x_1\times{x_2}\times{x_3}\times{...}\times{x_i} yi=x1×x2×x3×...×xi

用法：

python 复制代码

import torch
a = torch.Tensor([[1, 2, 3, 4, 5], [1, 2, 3, 4, 5]])
r0 = torch.cumprod(a, dim=0)
print(r0)
# tensor([[ 1.,  2.,  3.,  4.,  5.],
#         [ 1.,  4.,  9., 16., 25.]])

r1 = torch.cumprod(a, dim=1)
print(r1)
# tensor([[  1.,   2.,   6.,  24., 120.],
#         [  1.,   2.,   6.,  24., 120.]])

我们自习观察r0和r1的区别，在不同维度上进行累乘。更重要的是，每个阶段乘法结果都保存下来了，比如 1 × 2 × 3 × 4 × 5 1\times2\times3\times4\times5 1×2×3×4×5结果等于120，但前四步的结果1，2，6，24都保存下来了。这个计算刚好可以用来进行体渲染。