基于索引向量的张量缩减

【问题标题】：Tensor reduction based off index vector基于索引向量的张量缩减
【发布时间】：2019-04-26 14:46:51
【问题描述】：

例如，我有 2 个张量：A = [1;2;3;4;5;6;7] 和 B = [2;3;2]。这个想法是我想根据 B 减少 A - 这样 B 的值表示如何对 A 的值求和 - 这样B = [2;3;2] 意味着减少的 A 应该是前 2 个值、下一个 3 和最后 2 个值的总和： A' = [(1+2);(3+4+5);(6+7)]。很明显，B 的总和应始终等于 A 的长度。我试图尽可能有效地做到这一点 - 最好是 pytorch/python 中包含的特定函数或矩阵运算。谢谢！

【问题讨论】：

标签： python machine-learning pytorch

【解决方案1】：

这里是解决方案。

首先，我们创建一个索引数组B_idx，其大小与A 相同。
然后，使用index_add_，根据索引B_idx 累加（添加）A 中的所有元素。

A = torch.arange(1, 8) 
B = torch.tensor([2, 3, 2])

B_idx = [idx.repeat(times) for idx, times in zip(torch.arange(len(B)), B)]
B_idx = torch.cat(B_idx) # tensor([0, 0, 1, 1, 1, 2, 2])

A_sum = torch.zeros_like(B)
A_sum.index_add_(dim=0, index=B_idx, source=A)
print(A_sum) # tensor([ 3, 12, 13])

【讨论】：