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

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¹¹个元素，接近1TB，远远超出我的内存容量。

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

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

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

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

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

Answer 2

在 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 序的，则需要内存副本。但是内存访问是对齐的，这提供了很高的加速，尤其是在大型阵列上。它也很容易并行化。

计时

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

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

Efficiently avoid large intermediate array in computation?

python

numpy

numpy-ndarray