【记录】torch-discounted-cumsum工具
-
工具: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')