Python 中 3D 数组值的原始求和与 Numba JIT 基准测试

此函数汇总给定 3D 数组中的所有值：

def primitive_pixel_sum(frame):
    result = 0.0
    for x in range(frame.shape[0]):
        for y in range(frame.shape[1]):
            for z in range(frame.shape[2]):
                result += frame[x,y,z]
    return result

而以下函数是完全相同的算法，但使用 numba.jit：

import numba

@numba.jit
def numba_pixel_sum(frame):
    result = 0.0
    for x in range(frame.shape[0]):
        for y in range(frame.shape[1]):
            for z in range(frame.shape[2]):
                result += frame[x,y,z]
    return result

我们可以在 Jupyter 中使用以下命令对它们进行基准测试

%%timeit
primitive_pixel_sum(frame)

和

%%timeit
numba_pixel_sum(frame)

分别。

结果

We tested this with a r和om camera image taken from OpenCV of shape (480, 640, 3)

primitive_pixel_sum()：

1.78 s ± 253 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

numba_pixel_sum()：

4.06 ms ± 413 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

从这些结果应该清楚，numba 版本比原始版本快 438 倍。

注意，使用 numba.jit 编译复杂函数时，编译可能需要数毫秒甚至数秒 - 可能比简单 Python 函数花费的时间更长。

Since it’s so simple to use Numba, my recommendation is to just try it out for every function you suspect will eat up a lot of CPU time. Over time you will be able to develop an intuition for which functions it’s worth to use Numba 和 which functions won’t work at all or if it will be slower overall than just using Python.

记住，通常你也可以使用 NumPy 函数来实现相同的结果。在我们的示例中，你可以使用以下命令实现相同的功能

np.sum(frame)

这甚至比 Numba 更快：

%%timeit
np.sum(frame)

结果：

2.5 ms ± 7.17 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)