python - 使用 numba 加快拍摄速度？

Question

是否可以使用 numba 加速 np.take？

这是我的尝试，但速度要慢得多。我不能使用 nopython 模式，因为它不喜欢 np.empty_like 命令。

import numba
import numpy as np
from timer import Timer    

def take( x, indices ):
    result = np.empty_like( indices, dtype=x.dtype )    
    for i in range( len( indices ) ):
        result[i] = x[ indices[ i ] ]
    return result

jtake = numba.jit("f4[:](f4[:],i4[:])" )( take )

if __name__=='__main__':

    N = 100000
    m = 100
    idx = np.random.random_integers( 0, N, m )
    x = np.random.randn( N )

    num_tests=10000

    with Timer( 'take' ):
        for i in range( num_tests ):    
            r0 = take( x, idx )

    with Timer( 'Numba take' ):
        for i in range( num_tests ):    
            r1 = jtake( x, idx )                

    with Timer( 'Numpy.take' ):
        for i in range( num_tests ):
            r2 = x.take( idx )

有结果:

Beginning take
take took 2.46 seconds
Beginning Numba take
Numba take took 1.11 seconds
Beginning Numpy.take
Numpy.take took 0.04 seconds

Answer 1

答案是否定的。

Numba 不会作用于编译函数，例如 np.take() 或使用花式索引的数组方法，这是您的算法的基础。 Numba 作用于代码的解释部分。

您的 take() 函数可能比 NumPy 的开销更大，并且 Numba 改进了 for 循环(已解释)。

使用 NumPy > 1.9，您的代码应该更接近 NumPy 的做法，因为您的算法基于花式索引和 they improved the fancy indexing efficiency到 np.take() 的级别。