python - 分配给数组时 Numba 很慢？

Question

Numba 似乎是加速数字代码执行的绝佳解决方案。但是，当对数组进行赋值时，Numba 似乎比标准 Python 代码慢。考虑这个例子比较四个备选方案，使用/不使用 Numba，写入数组/标量:

(计算有意保持非常简单，以关注问题，即对标量的赋值与对数组单元格的赋值)

@autojit
def fast_sum_arr(arr):
    z = arr.copy()
    M = len(arr)
    for i in range(M):
        z[i] += arr[i]

    return z

def sum_arr(arr):
    z = arr.copy()
    M = len(arr)
    for i in range(M):
        z[i] += arr[i]

    return z

@autojit
def fast_sum_sclr(arr):
    z = 0
    M = len(arr)
    for i in range(M):
        z += arr[i]

    return z

def sum_sclr(arr):
    z = 0
    M = len(arr)
    for i in range(M):
        z += arr[i]

    return z

使用 IPython 的 %timeit 来评估我得到的四个备选方案:

In [125]: %timeit fast_sum_arr(arr)
100 loops, best of 3: 10.8 ms per loop

In [126]: %timeit sum_arr(arr)
100 loops, best of 3: 4.11 ms per loop

In [127]: %timeit fast_sum_sclr(arr)
100000 loops, best of 3: 10 us per loop

In [128]: %timeit sum_sclr(arr)
100 loops, best of 3: 2.93 ms per loop

sum_arr，不是用 Numba 编译的，比用 Numba 编译的 fast_sum_arr 快两倍多。另一方面，使用 Numba 编译的 fast_sum_sclr 比未使用 Numba 编译的 sum_sclr 快两个数量级以上。

所以 Numba 在加速 sum_sclr 的任务上表现得非常好，但实际上会使 sum_arr 执行得更慢。 sum_sclr 和 sum_arr 之间的唯一区别是前者分配给标量，而后者分配给数组单元格。

不知道有没有关系，最近在博客上看了下面的http://www.phi-node.com/ :

“事实证明，当 Numba 遇到它不直接支持的任何构造时，它会切换到(非常)慢的代码路径。”

博客作者使用 if 语句而不是 Python 的 max() 使 Numba 执行得更快。

对此有何见解？

谢谢，

全能

Answer 1

这里慢的是 arr.copy() 函数，而不是对数组的写访问。证明:

# -*- coding: utf-8 -*-
from numba import autojit
from Timer import Timer
import numpy as np

@autojit
def fast_sum_arr(arr, z):
    #z = arr.copy()
    M = len(arr)
    for i in range(M):
        z[i] += arr[i]

    return z

def sum_arr(arr, z):
    #z = arr.copy()
    M = len(arr)
    for i in range(M):
        z[i] += arr[i]

    return z

@autojit
def fast_sum_sclr(arr):
    z = 0
    M = len(arr)
    for i in range(M):
        z += arr[i]

    return z

def sum_sclr(arr):
    z = 0
    M = len(arr)
    for i in range(M):
        z += arr[i]

    return z

if __name__ == '__main__':
    vec1 = np.ones(1000)
    z = vec1.copy()
    with Timer() as t0:
        for i in range(10000):
            pass
    print "time for empty loop ", t0.secs
    print
    with Timer() as t1:
        for i in range(10000):
            sum_arr(vec1, z)
    print "time for sum_arr  [µs]:   ", (t1.secs-t0.secs)  / 10000 * 1e6
    with Timer() as t1:
        for i in range(10000):
            fast_sum_arr(vec1, z)
    print "time for fast_sum_arr  [µs]:   ", (t1.secs-t0.secs)  / 10000 * 1e6
    with Timer() as t1:
        for i in range(10000):
            sum_sclr(vec1)
    print "time for sum_arr  [µs]:   ", (t1.secs-t0.secs)  / 10000 * 1e6
    with Timer() as t1:
        for i in range(10000):
            fast_sum_sclr(vec1)
    print "time for fast_sum_arr  [µs]:   ", (t1.secs-t0.secs)  / 10000 * 1e6

"""
time for empty loop  0.000312089920044

time for sum_arr       [µs]:    432.02688694
time for fast_sum_arr  [µs]:      7.43598937988
time for sum_arr       [µs]:    284.574580193
time for fast_sum_arr  [µs]:      5.74610233307
"""