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

如何解决在计算中有效避免大型中间数组？

我有两个数组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 是否有工具可以提高效率（矢量化）？

解决方法

在 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],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)

,

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

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

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