python: 如何使基于 numba 的 for 循环更快

Question

我有一个for循环，它花费了很多时间。我想用numba模块来加速。

我的环境是：

win 10
python 3.7.5
anaconda 4.8.3
numpy 0.19.2
numba 0.46.0

原码为：

def computePoints(dxFullCurve, rows, columns, direction, relativeOffset, cprSpacing):
    points = []
    for row in range(rows):
        p = dxFullCurve[row, :]
        for col in range(columns):
            cprP = p.copy()
            cprP = cprP + direction * (col - columns / 2 - relativeOffset[row]) * cprSpacing
            points.append(cprP)
    return points


if __name__ == '__main__':

    dxFullCurve = np.random.random(size=[500, 3])
    direction = np.array([1, 0, 0])
    rows = 500
    columns = 500
    relativeOffset = np.random.random(size=500)
    cprSpacing = 0.1
    import time
    t1 = time.time()
    for i in range(100):
        computePoints(dxFullCurve, rows, columns, direction, relativeOffset, cprSpacing)
    t2 = time.time()
    print('time: ', (t2-t1)/100)

打印时间为：0.8

然后，我用numba来加速，代码是：

@nb.jit()
def computePoints(dxFullCurve, rows, columns, direction, relativeOffset, cprSpacing):
    points = []
    for row in range(rows):
        p = dxFullCurve[row, :]
        for col in range(columns):
            cprP = p.copy()
            cprP = cprP + direction * (col - columns / 2 - relativeOffset[row]) * cprSpacing
            points.append(cprP)
    return points

现在，时间是：0.177。 numba 真的加快了速度。然而，它只加速了 4 倍。有什么方法可以让它更快吗？

然后，我尝试了 numba parallel 如下：

@nb.jit(nopython=True, parallel=True)
def computePoints(dxFullCurve, rows, columns, direction, relativeOffset, cprSpacing):
    points = []
    for row in range(rows):
        p = dxFullCurve[row, :]
        for col in range(columns):
            cprP = p.copy()
            cprP = cprP + direction * (col - columns / 2 - relativeOffset[row]) * cprSpacing
            points.append(cprP)
    return points

但是，成本时间是：0.903。难以置信，它甚至比非 numba 代码花费更多的时间。

我只想知道：有什么方法可以让我的for循环更快吗？

Answer 1

您至少有两件事会减慢您的代码速度：

p.copy() 是不必要的。只需删除行 cprP = p.copy() 并更改为 cprP = p + direction * ....
Numba 仅适用于数组，但您可以在 list 中累积结果。据我所知，你所有的个人点都是 (3,) 形状的数组，你有 rows*columns 个。在下面的代码中我pre-allocate points作为一个数组然后在循环中填充它。

@nb.jit
def computePoints(dxFullCurve, rows, columns, direction, relativeOffset, cprSpacing):
    points = np.empty((rows*columns, 3))
    index = 0
    for row in range(rows):
        p = dxFullCurve[row, :]
        for col in range(columns):
            cprP = p + direction * (col - columns / 2 - relativeOffset[row]) * cprSpacing
            points[index, :] = cprP
            index += 1
    return points

这两项更改使我的机器上的速度额外提高了 8 倍。

Answer 2

如何优化 Numba 函数

这是对@jmd_dk回答的较长评论。缺少一些重要的点，进一步加快了计算速度。

显式内循环 -> 更容易为编译器优化在这种情况下，很可能展开内循环
移除循环迭代之间不必要的依赖（索引变量）。这使得并行化成为可能，并且通常是一个好主意，因为 CPU 核心具有多个计算单元。
parallel=True 启用并行化。这仅在运行时足够大时才有用。如果一个函数只需要几微秒，请不要这样做。
fastmath=True -> 代数变化是允许的，在数值上这可能会对结果产生影响，程序员必须决定是否可以。
error_model='numpy'-> 关闭除以零的检查，只有在真正的除法上才真正需要这个可以优化到 *0.5
cache=True 如果使用相同数据类型的输入调用该函数，则该函数只需在您重新启动解释器时从缓存中加载。如果您有更复杂的功能，这将特别有用
尽可能避免列表（jmd_dk
use assert statements: 这不仅是为了安全（没有边界检查），而且还告知编译器确切的内存布局。在不知道这一点的情况下，编译器通常无法确定 SIMD-vectorization 是否有益。你在这里看不到太多加速，因为这只是复制数据，中间有一些微不足道的乘法，但其他功能的加速可能是实质性的
进一步优化：尽可能避免内存分配。这是迄今为止整个功能中成本最高的部分（至少如果实施了前面的几点）。

例子

@nb.njit(fastmath=True,error_model="numpy",parallel=True)
def computePoints_nb_2(dxFullCurve, rows, columns, direction, relativeOffset, cprSpacing):
    assert dxFullCurve.shape[1]==3
    assert direction.shape[0]==3
          
    points = np.empty((rows*columns, 3))
    for row in nb.prange(rows):
        for col in range(columns):
            for i in range(3):
                points[row*columns+col, i] = dxFullCurve[row, i] + direction[i] * (col - columns / 2 - relativeOffset[row]) * cprSpacing
    return points

如果可以避免内存分配。

@nb.njit(fastmath=True,error_model="numpy",parallel=True)
def computePoints_nb_2_pre(dxFullCurve, rows, columns, direction, relativeOffset, cprSpacing,points):
    assert dxFullCurve.shape[1]==3
    assert direction.shape[0]==3
    assert points.shape[1]==3

    for row in nb.prange(rows):
        for col in range(columns):
            for i in range(3):
                points[row*columns+col, i] = dxFullCurve[row, i] + direction[i] * (col - columns / 2 - relativeOffset[row]) * cprSpacing
    return points

计时

#Implementation of jmd_dk
%timeit computePoints_nb_1(dxFullCurve, rows, columns, direction, relativeOffset, cprSpacing)
#23.2 ms ± 213 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit computePoints_nb_2(dxFullCurve, rows, columns, direction, relativeOffset, cprSpacing)
#1.54 ms ± 61.5 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit computePoints_nb_2_pre(dxFullCurve, rows, columns, direction, relativeOffset, cprSpacing,points)
#122 µs ± 4.1 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

python: 如何使基于 numba 的 for 循环更快

python: how to make the numba based for loop faster

python

numba

如何优化 Numba 函数