python - 提高比较算法 np.packbits(A==A[ :, None], axis=1) 的性能

Question

我正在寻求内存优化 np.packbits(A==A[:, None], axis=1)，其中 A 是密集的整数数组长度 n。 A==A[:, None] 需要很大的 n，因为生成的 boolean 数组存储效率低下，每个 boolean 值占用 1 个字节。

我编写了以下脚本来实现相同的结果，同时一次打包一个部分。然而，它慢了大约 3 倍，所以我正在寻找加速它的方法。或者，另一种方法是内存开销较小的更好算法。

注意:这是我之前提出的问题的后续问题； Comparing numpy array with itself by element efficiently .

下面用于基准测试的可重现代码。

import numpy as np
from numba import jit

@jit(nopython=True)
def bool2int(x):
    y = 0
    for i, j in enumerate(x):
        if j: y += int(j)<<(7-i)
    return y

@jit(nopython=True)
def compare_elementwise(arr, result, section):
    n = len(arr)
    for row in range(n):
        for col in range(n):

            section[col%8] = arr[row] == arr[col]

            if ((col + 1) % 8 == 0) or (col == (n-1)):
                result[row, col // 8] = bool2int(section)
                section[:] = 0

    return result

n = 10000
A = np.random.randint(0, 1000, n)

result_arr = np.zeros((n, n // 8 if n % 8 == 0 else n // 8 + 1)).astype(np.uint8)
selection_arr = np.zeros(8).astype(np.uint8)

# memory efficient version, but slow
packed = compare_elementwise(A, result_arr, selection_arr)

# memory inefficient version, but fast
packed2 = np.packbits(A == A[:, None], axis=1)

assert (packed == packed2).all()

%timeit compare_elementwise(A, result_arr, selection_arr)  # 1.6 seconds
%timeit np.packbits(A == A[:, None], axis=1)  # 0.460 second

Answer 1

这是一个比 numpy 快 3 倍的解决方案(a.size 必须是 8 的倍数；见下文):

@nb.njit
def comp(a):
    res=np.zeros((a.size,a.size//8),np.uint8)
    for i,x in enumerate(a):
        for j,y in enumerate(a):
            if x==y: res[i,j//8] |= 128 >> j%8 
    return res

这是有效的，因为数组被扫描了一次，而你做了很多次， 并且几乎所有条款都是空的。

In [122]: %timeit np.packbits(A == A[:, None], axis=1)
389 ms ± 57.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [123]: %timeit comp(A)
123 ms ± 24.4 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

如果a.size%8 > 0，找回信息的代价会更高。在这种情况下，最好的方法是用一些(在 range(7) 中)零填充初始数组。

为了完整起见，可以这样填充:

if A.size % 8 != 0: A = np.pad(A, (0, 8 - A.size % 8), 'constant', constant_values=0)