使用 CuPy 或 NumPy 高效计算分区总和

Posted 2023-02-16

【中文标题】使用 CuPy 或 NumPy 高效计算分区总和

【英文标题】：Calculate partitioned sum efficiently with CuPy or NumPy

【发布时间】：2020-07-10 14:36:30

【问题描述】：

我有一个很长的数组*，长度为 L（我们称之为 values），我想对其求和，以及一个相同长度的排序一维数组 L，其中包含 N 整数对原始数组进行分区——我们称这个数组为labels。

我目前正在做的是这个（module 是 cupy 或 numpy）：

result = module.empty(N)
for i in range(N):
    result[i] = values[labels == i].sum()

但这不是最有效的方法（应该可以摆脱for 循环，但如何？）。由于labels 已排序，我可以轻松确定断点并将这些索引用作开始/停止点，但我看不出这如何解决for 循环问题。

请注意，如果可能的话，我想避免创建大小为 NxL 的数组，因为 L 非常大。

我在 cupy 工作，但也欢迎任何 numpy 解决方案，并且可能会被移植。在 cupy 中，ReductionKernel 似乎就是这种情况，但我不太明白该怎么做。

---

* 在我的例子中，values 是一维的，但我认为解决方案不会依赖于此

【参考方案1】：

您正在描述 groupby sum 聚合。您可以为此编写一个 CuPy RawKernel，但使用在 GPU 数据帧库 cuDF 中实现的现有 groupby 聚合会非常更容易。它们可以互操作，而无需您复制数据。如果您在生成的 cuDF 系列上调用 .values，它将为您提供一个 CuPy 数组。

如果你回到 CPU，你可以用 pandas 做同样的事情。

import cupy as cp
import pandas as pd

N = 100
values = cp.random.randint(0, N, 1000)
labels = cp.sort(cp.random.randint(0, N, 1000))

L = len(values)
result = cp.empty(L)
for i in range(N):
    result[i] = values[labels == i].sum()
    
result[:5]
array([547., 454., 402., 601., 668.])

import cudf

df = cudf.DataFrame("values": values, "labels": labels)
df.groupby(["labels"])["values"].sum().values[:5]
array([547, 454, 402, 601, 668])

【讨论】：

谢谢，我一定会研究 cuDF 并分析/时间不同版本的相互对比，这可能是一个很好的解决方案。但是我认为您在发布的第一个 sn-p 中忘记使用 pandas 了？似乎与我目前正在做的事情相同（即使用for）。

糟糕，我明白了你想要展示的内容，抱歉，忘了我说的关于熊猫的内容。

另外，对于垃圾邮件，我很抱歉，但我刚刚意识到我的 results 数组的长度应该是 N，而不是 L。我想你可能也对此感到恼火。

【参考方案2】：

这是一个解决方案，它使用N x <max partition size in labels> 数组而不是N x L 数组（如果不同分区之间的差异不太大，则该数组不应该很大）：

将数组调整为二维数组，每行都有分区。由于行的长度等于最大分区的大小，因此用零填充不可用的值（因为它不会影响任何总和）。这使用了@Divakar 给出的here 的解决方案。

def jagged_to_regular(a, parts):
    lens = np.ediff1d(parts,to_begin=parts[0])
    mask = lens[:,None]>np.arange(lens.max())
    out = np.zeros(mask.shape, dtype=a.dtype)
    out[mask] = a
    return out

parts_stack = jagged_to_regular(values, labels)

沿轴 1 求和：

result = np.sum(parts_stack, axis = 1)

如果您想要一个 CuPy 实现，jagged_to_regular 中没有直接的 CuPy alternative 到 numpy.ediff1d。在这种情况下，您可以将语句替换为 numpy.diff，如下所示：

lens = np.insert(np.diff(parts), 0, parts[0])

然后继续使用 CuPy 作为 numpy 的替代品。

【讨论】：

有趣，谢谢。我还将针对 Nick Becker 的解决方案对此进行测试，看看它们各自的表现如何。根据我的经验，我怀疑这里的布尔屏蔽操作可能有点内存效率低下。