为什么numpy.where比其他选择快得多

Question

我正在尝试加速以下代码：

import time
import numpy as np
np.random.seed(10)
b=np.random.rand(10000,1000)
def f(a=1):
    tott=0
    for _ in range(a):
        q=np.array(b)
        t1 = time.time()
        for i in range(len(q)):
            for j in range(len(q[0])):
                if q[i][j]>0.5:
                    q[i][j]=1
                else:
                    q[i][j]=-1
        t2=time.time()
        tott+=t2-t1
    print(tott/a)

正如你所看到的，func主要是关于在双循环中迭代。因此，我尝试使用np.nditer、np.vectorize和map来代替它。如果提供了一些加速(比如4-5倍，除了np.nditer)，但是！使用np.where(q>0.5,1,-1)时，加速比几乎是100倍。我怎样才能像np.where一样快速地迭代numpy数组呢？为什么它要快这么多？

Answer 1

这是因为numpy的核心是用C实现的，你基本上是在比较C和Python的速度。

如果想要利用numpy的速度优势，就应该在Python代码中进行尽可能少的调用。如果使用Python循环，即使只在该循环中使用numpy函数，也已经失败了。使用numpy提供的更高级别的函数(这就是为什么他们提供了这么多特殊函数)。在内部，它将使用效率更高的(C-)循环

你可以自己用C(带循环)实现一个函数，然后从Python中调用它。这应该会给出类似的速度。

Answer 2

要回答这个问题，您可以通过使用numba库来获得相同的速度(100倍加速)：

from numba import njit

def f(b):
    q = np.zeros_like(b)

    for i in range(b.shape[0]):
        for j in range(b.shape[1]):
            if q[i][j] > 0.5:
                q[i][j] = 1
            else:
                q[i][j] = -1

    return q

@njit
def f_jit(b):
    q = np.zeros_like(b)

    for i in range(b.shape[0]):
        for j in range(b.shape[1]):
            if q[i][j] > 0.5:
                q[i][j] = 1
            else:
                q[i][j] = -1

    return q

比较速度：

朴素的Python

%timeit f(b)
592 ms ± 5.72 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Numba (使用LLVM ~C速度进行即时编译)

%timeit f_jit(b)
5.97 ms ± 105 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)