嗯,这基本上是做
np.bincount与做
1D阵列。但是,我们需要在每行上迭代使用它(简单地考虑一下)。为了使其向量化,我们可以使每行偏移最大数量。想法是每行具有不同的bin,以使它们不受编号相同的其他行元素的影响。
因此,实施将是-
# Vectorized solutiondef 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]: aOut[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并行循环和Cython的prange。prange
在较大的数据集上表现良好,这可能是由于设置并行工作所需的开销。
因此,通过这两项新的调整以及
njit针对非Python模式的调整,我们将获得三种变体-
# Numba solutionsdef 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@njitdef 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 solutiondef 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 loop10000 loops, best of 3: 36.7 µs per loop10000 loops, best of 3: 22.6 µs per loop10000 loops, best of 3: 22.7 µs per loop10000 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 loop100 loops, best of 3: 3.54 ms per loop1000 loops, best of 3: 1.83 ms per loop100 loops, best of 3: 1.78 ms per loop1000 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 loop100 loops, best of 3: 4.86 ms per loop100 loops, best of 3: 3.21 ms per loop100 loops, best of 3: 3.18 ms per loop100 loops, best of 3: 2.45 ms per loop
看起来这些
numba变体的效果非常好。从三个变体中选择一个将取决于输入数组的形状参数,并在某种程度上取决于其中的唯一元素的数量。



