【记录】torch-discounted-cumsum工具

[183****0229]

工具：torch-discounted-cumsum
顾名思义：不连续 求和操作。
地址：https://github.com/toshas/torch-discounted-cumsum

该存储库实现了一种高效的并行算法，用于计算不连续的累积总和。不连续的cumsum操作经常出现在与时间序列相关的数据科学领域，包括强化学习 (RL)。
传统的顺序算法在循环中执行输出元素的计算。对于 size 的输入 N，它需要O(N)操作并需要O(N)时间步骤才能完成。

并行算法的特点：

• 输入大小的速度对数
• 比顺序算法更好的数值精度

包的特点：

• CPU: sequential algorithm in C++
• GPU: parallel algorithm in CUDA
• Gradients computation for input and gamma
• Batch support for input and gamma
• Both left and right directions of summation supported
• PyTorch bindings

安装

pip install torch-discounted-cumsum

使用

import torch
from torch_discounted_cumsum import discounted_cumsum_right

N = 8
gamma = 0.99
x = torch.ones(1, N).cuda()
y = discounted_cumsum_right(x, gamma)

print(y)
tensor([[7.7255, 6.7935, 5.8520, 4.9010, 3.9404, 2.9701, 1.9900, 1.0000]],
       device='cuda:0')