python - 每行 Bin 元素 - NumPy 的矢量化 2D Bincount

Question

我有一个带有整数值的 NumPy 数组。矩阵的值范围从 0 到矩阵中的最大元素(换句话说，矩阵中呈现从 0 到最大数据元素的所有数字)。我需要构建有效的(有效意味着快速完全矢量化解决方案)来搜索每行中的元素数量并根据矩阵值对它们进行编码。

我找不到类似的问题，或者以某种方式帮助解决这个问题的问题。

所以如果我在输入中有这个数据:

# shape is (N0=4, m0=4) 
1   1   0   4
2   4   2   1
1   2   3   5
4   4   4   1

期望的输出是:

# shape(N=N0, m=data.max()+1):
1   2   0   0   1   0
0   1   2   0   1   0
0   1   1   1   0   1
0   1   0   0   3   0

我知道如何解决这个问题，只需逐一迭代计算每行 data 中的唯一值，然后考虑 data 中所有可能的值来组合结果大批。

虽然使用 NumPy 进行向量化，但关键问题是逐一搜索每个数字很慢，并且假设存在大量唯一数字，这不是有效的解决方案。一般来说，N 和唯一数字计数都相当大(顺便说一下，N 似乎比唯一数字计数要大)。

有人有很棒的想法吗？)

Answer 1

这基本上就是 np.bincount 的作用。与1D 数组一样。但是，我们需要在每一行迭代地使用它(简单地思考一下)。为了使其矢量化，我们可以将每行偏移该最大数。这个想法是为每行设置不同的 bin，这样它们就不会受到具有相同编号的其他行元素的影响。

因此，实现将是 -

# Vectorized solution
def bincount2D_vectorized(a):    
    N = a.max()+1
    a_offs = a + np.arange(a.shape[0])[:,None]*N
    return np.bincount(a_offs.ravel(), minlength=a.shape[0]*N).reshape(-1,N)

示例运行 -

In [189]: a
Out[189]: 
array([[1, 1, 0, 4],
       [2, 4, 2, 1],
       [1, 2, 3, 5],
       [4, 4, 4, 1]])

In [190]: bincount2D_vectorized(a)
Out[190]: 
array([[1, 2, 0, 0, 1, 0],
       [0, 1, 2, 0, 1, 0],
       [0, 1, 1, 1, 0, 1],
       [0, 1, 0, 0, 3, 0]])

Numba 调整

我们可以引入numba以进一步加快速度。现在，numba 允许进行一些调整。

• 首先，它允许 JIT 编译。

• 此外，最近他们推出了实验性的 parallel自动并行化已知具有并行语义的函数中的操作。

• 最后的调整是使用 prange作为 range 的替代品。文档指出，这会并行运行循环，类似于 OpenMP 并行 for 循环和 Cython 的 prange。 prange 在较大的数据集上表现良好，这可能是因为设置并行工作所需的开销。

因此，通过这两项新的调整以及 njit对于非 Python 模式，我们将有三种变体 -

# Numba solutions
def bincount2D_numba(a, use_parallel=False, use_prange=False):
    N = a.max()+1
    m,n = a.shape
    out = np.zeros((m,N),dtype=int)

    # Choose fucntion based on args
    func = bincount2D_numba_func0
    if use_parallel:
        if use_prange:
            func = bincount2D_numba_func2
        else:
            func = bincount2D_numba_func1
    # Run chosen function on input data and output
    func(a, out, m, n)
    return out

@njit
def bincount2D_numba_func0(a, out, m, n):
    for i in range(m):
        for j in range(n):
            out[i,a[i,j]] += 1

@njit(parallel=True)
def bincount2D_numba_func1(a, out, m, n):
    for i in range(m):
        for j in range(n):
            out[i,a[i,j]] += 1

@njit(parallel=True)
def bincount2D_numba_func2(a, out, m, n):
    for i in prange(m):
        for j in prange(n):
            out[i,a[i,j]] += 1

为了完整性和稍后测试，循环版本将是 -

# Loopy solution
def bincount2D_loopy(a):
    N = a.max()+1
    m,n = a.shape
    out = np.zeros((m,N),dtype=int)
    for i in range(m):
        out[i] = np.bincount(a[i], minlength=N)
    return out 

运行时测试

案例#1:

In [312]: a = np.random.randint(0,100,(100,100))

In [313]: %timeit bincount2D_loopy(a)
     ...: %timeit bincount2D_vectorized(a)
     ...: %timeit bincount2D_numba(a, use_parallel=False, use_prange=False)
     ...: %timeit bincount2D_numba(a, use_parallel=True, use_prange=False)
     ...: %timeit bincount2D_numba(a, use_parallel=True, use_prange=True)
10000 loops, best of 3: 115 µs per loop
10000 loops, best of 3: 36.7 µs per loop
10000 loops, best of 3: 22.6 µs per loop
10000 loops, best of 3: 22.7 µs per loop
10000 loops, best of 3: 39.9 µs per loop

案例#2:

In [316]: a = np.random.randint(0,100,(1000,1000))

In [317]: %timeit bincount2D_loopy(a)
     ...: %timeit bincount2D_vectorized(a)
     ...: %timeit bincount2D_numba(a, use_parallel=False, use_prange=False)
     ...: %timeit bincount2D_numba(a, use_parallel=True, use_prange=False)
     ...: %timeit bincount2D_numba(a, use_parallel=True, use_prange=True)
100 loops, best of 3: 2.97 ms per loop
100 loops, best of 3: 3.54 ms per loop
1000 loops, best of 3: 1.83 ms per loop
100 loops, best of 3: 1.78 ms per loop
1000 loops, best of 3: 1.4 ms per loop

案例#3:

In [318]: a = np.random.randint(0,1000,(1000,1000))

In [319]: %timeit bincount2D_loopy(a)
     ...: %timeit bincount2D_vectorized(a)
     ...: %timeit bincount2D_numba(a, use_parallel=False, use_prange=False)
     ...: %timeit bincount2D_numba(a, use_parallel=True, use_prange=False)
     ...: %timeit bincount2D_numba(a, use_parallel=True, use_prange=True)
100 loops, best of 3: 4.01 ms per loop
100 loops, best of 3: 4.86 ms per loop
100 loops, best of 3: 3.21 ms per loop
100 loops, best of 3: 3.18 ms per loop
100 loops, best of 3: 2.45 ms per loop

看起来 numba 变体表现得非常好。从三种变体中选择一种取决于输入数组形状参数，并在某种程度上取决于其中唯一元素的数量。