python - 在计算中有效避免大型中间数组？

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 中增加 block 大小可能会有所帮助，但在编译代码中实现算法也非常简单。我在这里使用了 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 序的，则需要内存副本。但是内存访问是对齐的，这提供了很高的加速，尤其是在大型阵列上。它也很容易并行化。

时间

#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)