我想将一个长向量分成较小的不相等片段,对每个片段进行求和,然后将结果收集到一个新向量中。 我需要在pytorch中执行此操作,但我也有兴趣查看如何使用numpy完成此操作。
这可以通过拆分向量轻松实现。
sizes = [3, 7, 5, 9]
X = torch.ones(sum(sizes))
Y = torch.tensor([s.sum() for s in torch.split(X, sizes)])
或np.ones和np.split。
有没有更有效的方法?
灵感来自第一条评论:
indices = np.cumsum([0]+sizes)[:-1]
Y = np.add.reduceat(X, indices.tolist())
将其解析为numpy。我仍在用pytorch寻找解决方案。
答案 0 :(得分:0)
index_add_
是你的朋友!
# inputs
sizes = torch.tensor([3, 7, 5, 9], dtype=torch.long)
x = torch.ones(sizes.sum())
# prepare an index vector for summation (what elements of x are summed to each element of y)
ind = torch.zeros(sizes.sum(), dtype=torch.long)
ind[torch.cumsum(sizes, dim=0)[:-1]] = 1
ind = torch.cumsum(ind, dim=0)
# prepare the output
y = torch.zeros(len(sizes))
# do the actual summation
y.index_add_(0, ind, x)