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 并行，如下所示：

@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

如何优化 Numba 函数

\n

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

\n

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

\n

例子

\n

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

\n

如果可以避免内存分配。

\n

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

\n

时间安排

\n

#Implementation of jmd_dk\n%timeit computePoints_nb_1(dxFullCurve, rows, columns, direction, relativeOffset, cprSpacing)\n#23.2 ms \xc2\xb1 213 \xc2\xb5s per loop (mean \xc2\xb1 std. dev. of 7 runs, 1 loop each)\n\n%timeit computePoints_nb_2(dxFullCurve, rows, columns, direction, relativeOffset, cprSpacing)\n#1.54 ms \xc2\xb1 61.5 \xc2\xb5s per loop (mean \xc2\xb1 std. dev. of 7 runs, 1 loop each)\n\n%timeit computePoints_nb_2_pre(dxFullCurve, rows, columns, direction, relativeOffset, cprSpacing,points)\n#122 \xc2\xb5s \xc2\xb1 4.1 \xc2\xb5s per loop (mean \xc2\xb1 std. dev. of 7 runs, 10000 loops each)\n

\n