有效避免计算中的大型中间数组？

Question

我有两个数组p 和alpha，我想<=-比较p 的每个元素与alpha 的每个元素，然后在第一个轴上聚合（计数）。我的代码：

s = np.sum(np.less_equal.outer(p, alpha), axis=0)

p 至少是一维的，但也可能是多维的，可以有 100 × 1000000 这样的维度。alpha 是一维的，通常有 100 到 1000 个元素。

问题是np.less_equal.outer 创建了一个中间数组，在最坏的情况下，它的大小可能是 100 × 1000000 × 1000 = 10¹¹ 个元素，接近 1 TB，远远超出我的记忆容量。

我的方法是沿第一个轴拆分计算：

s = np.zeros(shape=(*p.shape[1:], len(alpha)), dtype='int64')
for pr in p:
    s += np.less_equal.outer(pr, alpha)

这似乎可行，但我想知道 NumPy 是否有工具可以提高效率（矢量化）？

Answer 1

在 numpy 中增加块大小可能会有所帮助，但在编译后的代码中实现算法也非常简单。我在这里使用了 Numba，但进行了一些更改后，它应该也可以在 Cython 中使用。

第一次尝试

import numpy as np
import numba as nb
@nb.njit(fastmath=True)
def outer_cust(p,alpha):
    res=np.zeros((p.shape[1],alpha.shape[0]),dtype=np.int64)
    for i in range(p.shape[0]):
        for j in range(p.shape[1]):
            for k in range(alpha.shape[0]):
                if p[i,j]<=alpha[k]:
                    res[j,k]+=1
                else:
                    res[j,k]+=0
    return res

此方法不使用任何临时数组。然而，内存访问没有对齐（假设 C 连续数组）。因此性能不会很好，但已经比numpy算法快了。

第二次尝试

@nb.njit(fastmath=True,parallel=True,cache=True)
def outer_cust_2(p,alpha):
    #p should be fortran ordered to avoid memory copy
    p_T=np.ascontiguousarray(p.T)
    res=np.empty((p.shape[1],alpha.shape[0]),dtype=np.int64)
    for j in nb.prange(p_T.shape[0]):
        for k in range(alpha.shape[0]):
            acc=0
            for i in range(p_T.shape[1]):
                if p_T[j,i]<=alpha[k]:
                    acc+=1
            res[j,k]=acc
    return res

在这种方法中，如果输入是 c-ordered，则需要内存副本。但是内存访问是对齐的，这提供了很高的加速，尤其是在大型阵列上。它也很容易并行化。

时间

#Smaller arrays
p=np.random.rand(1000,10000)
alpha=np.random.rand(1000)
%timeit outer_cust_np(p,alpha)
#20 s ± 580 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit outer_cust(p,alpha)
#5.56 s ± 61.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit outer_cust_2(p,alpha)
#166 ms ± 1.19 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

#Maximum size
p=np.random.rand(1000,1000000)
alpha=np.random.rand(1000)
%timeit outer_cust_np(p,alpha)
#too slow to measure
%timeit outer_cust(p,alpha)
#too slow to measure
%timeit outer_cust_2(p,alpha)
#24.2 s ± 1.52 s per loop (mean ± std. dev. of 7 runs, 1 loop each)

Answer 2

在我看来，您正在为一些下游计算生成数据。如果特定的下游计算可以分部分完成，那么您可以逐列进行，将结果保存到磁盘，然后再合并结果。

另一种方法是在下游计算期间使用 p 和 alpha 上的索引操作按需获取这个大型矩阵的所需部分。

出于简单的原因，我更喜欢第二种方法。