时间戳

首页 > 解决方案 > 有效避免计算中的大型中间数组？

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

问题描述

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

标签： pythonnumpynumpy-ndarray

解决方案

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

推荐阅读