NumPy 数组行差异

Question

我有一个 NumPy 数组 vectors = np.random.randn(rows, cols)。我想根据其他一些稀疏数组 diffs 找到其行之间的差异和“2-hot”：在对应于 vectors 的第一行的列中包含 1 和a -1 对应第二行。也许举个例子更清楚：

diffs = np.array([[ 1,  0, -1],
                  [ 1, -1,  0]])

然后我可以通过 diffs @ vectors.

简单地计算行差异

不幸的是，对于 10_000x1000 的 diffs 和矢量 1000x15_000，这很慢。我可以通过使用 scipy.sparse 获得加速：sparse.csr_matrix(diffs) @ vectors，但即使这样也是 300ms.

可能这已经是最快的速度了，但我的某些部分认为使用矩阵乘法是否是完成这项任务的最明智的决定。

更重要的是，我之后需要取绝对值，所以我真的在做 np.abs(sparse.csr_matrix(diffs) @ vectors)，它添加了 ~ 200ms 总计 ~500ms.

Answer 1

I can compute the row differences by simply diffs @ vectors.

这是非常低效的。矩阵乘法在 O(n*m*k) 中运行 (n,m) 乘以 (m,k)。在您的情况下，每行只有两个值，您实际上不需要乘以 1 或 -1。您的问题可以在 O(n*k) 时间内计算出来（即快 m 倍）。

Unfortunately this is slow for diffs of 10_000x1000 and vectors 1000x15_000. I can get a speedup by using scipy.sparse.

问题是输入数据表示效率低下。当 diff 是一个大小为 (10_000,1000) 的数组时，使用比需要大 ~1000 倍的密集矩阵或未针对只有两个 [=51] 进行优化的稀疏矩阵都是不合理的=] 值（尤其是 1 和 -1）。您需要将 non-zeros 值的位置存储在形状为 (2,n) 的名为 sel_rows 的二维数组中，其中第一行包含 1 的位置，第二行包含 1 的位置diff 二维数组中的 -1。然后，您可以提取 vectors 的行，例如使用 vectors[sel_rows[0]]。您可以使用 vectors[sel_rows[0,:]] - vectors[sel_rows[1,:]] 执行最后的操作。这种方法应该比密集矩阵产品快得多，并且它可能比关于目标机器的稀疏矩阵产品快一点。

虽然上面的解决方案很简单，但它创建了多个不是 cache-friendly 的临时数组，因为您的输出数组应该采用 10_000 * 15_000 * 8 = 1.1 GiB（这相当巨大 ).您可以使用 Numba 删除临时数组，从而提高性能。可以使用多个线程来进一步提高性能。这是一个未经测试的代码：

import numba as nb

@nb.njit('(int_[:,::1], float64[:,::1])', parallel=True)
def compute(diffs, vectors):
    n, k = diffs.shape[0], vectors.shape[1]
    assert diffs.shape[1] == 2
    res = np.empty((n, k))
    for i in nb.prange(n):
        a, b = diffs[i]
        for j in range(k):
            # Compute nb.abs here if needed so to avoid 
            # creating new temporary arrays
            res[i, j] = vectors[a, j] - vectors[b, j]
    return res

上面的代码应该接近最优。它应该是 内存限制 并且能够 使内存带宽饱和 。请注意，在内存中写入这样 巨大的数组 需要一些时间以及读取（两次）输入数组。在 x86-64 平台上，基本实现应该将 4.4 GiB 数据 from/to 移动到 RAM。因此，在具有 20 GiB/s RAM 的主流 PC 上，这需要 220 毫秒。事实上，对于顺序实现，稀疏矩阵计算结果在实践中并没有那么糟糕。

如果这对您来说还不够，那么您可以使用 simple-precision floating-point 数字代替 double-precision（快两倍）。您还可以使用 low-level C/C++ 实现来减少使用的内存带宽（感谢 non-temporal 指令——快 ~30%）。没什么可做的了。

NumPy 数组行差异

NumPy array row differences

python

arrays

performance

numpy

scipy