numba 函数中的这个循环可以优化以运行得更快吗？

Question

在我的python代码中，我需要循环大约 2500 万次，我希望尽可能对其进行优化。循环内的操作非常简单。为了使代码高效，我使用了numba模块，这有很大帮助，但如果可能的话我想进一步优化代码。

这是一个完整的工作示例：

import numba as nb
import numpy as np
import time 
#######create some synthetic data for illustration purpose##################
size=5000
eps = 0.2
theta_c = 0.4
temp = np.ones(size)
neighbour = np.random.randint(size, size=(size, 3)) 
coschi = np.random.random_sample((size))
theta = np.random.random_sample((size))*np.pi/2
pwr = np.cos(theta)
###################end of dummy data##########################

###################-----main loop------###############
@nb.jit(fastmath=True)
def func(theta, pwr, neighbour, coschi, temp):
    for k in range(np.argmax(pwr), 5000*(pwr.size)): 
        n = k%pwr.size
        if (np.abs(theta[n]-np.pi/2.)<np.abs(theta_c)):
                adj = neighbour[n,1]
        else:
                adj = neighbour[n,0]
        psi_diff = np.abs(np.arccos(coschi[adj])-np.arccos(coschi[n]))
        temp5 = temp[adj]**5;
        e_temp = 1.- np.exp(-temp5*psi_diff/np.abs(eps))
        temp[n] = temp[adj] + (e_temp)/temp5*(pwr[n] - temp[adj]**4)
    return temp

#check time
time1 = time.time()
temp = func(theta, pwr, neighbour, coschi, temp)
print("Took: ", time.time()-time1, " seconds.")

这3.49 seconds在我的机器上进行。

为了某些模型拟合的目的，我需要运行这段代码几千次，因此即使是 1 秒的优化也意味着为我节省了数十个小时。

可以做些什么来进一步优化这段代码？

Answer 1

让我首先提出一些一般性评论：

\n\n

\n

• 如果您使用 numba 并且真正关心性能，您应该避免 numba 创建对象模式代码的任何可能性。这意味着您应该使用numba.njit(...)ornumba.jit(nopython=True, ...)代替numba.jit(...).

  \n\n

  这对您的情况没有影响，但它使意图更清晰，并且一旦（快速）nopython 模式不支持某些内容，就会抛出异常。

\n

• 你应该小心你的时间和方式。对 numba-jitted 函数（未提前编译）的第一次调用将包括编译成本。因此，您需要在计时之前运行一次以获得准确的计时。为了获得更准确的计时，您应该多次调用该函数。我喜欢%timeitJupyter Notebooks/Lab 中的 IPython 来粗略了解性能。

  \n\n

  所以我会使用：

  \n\n

  res1 = func(theta, pwr, neighbour, coschi, np.ones(size))\nres2 = # other approach\n\nnp.testing.assert_allclose(res1, res2)\n\n%timeit func(theta, pwr, neighbour, coschi, np.ones(size))\n%timeit # other approach\n

  \n\n

  这样，我将第一次调用（包括编译时间）与断言一起使用，以确保它确实产生（几乎）相同的输出，然后使用更强大的计时方法（与 相比time）对函数进行计时。

\n

\n\n

吊装np.arccos

\n\n

现在让我们从一些实际的性能优化开始：一个明显的优化是您可以提升一些“不变量”，例如 的np.arccos(coschi[...])计算频率比coschi. 您对每个元素进行大约 5000 次迭代，并且每个循环coschi执行两次！np.arccos因此，让我们计算arccos一次coschi并将其存储在中间数组中，以便可以在循环内访问它：

\n\n

res1 = func(theta, pwr, neighbour, coschi, np.ones(size))\nres2 = # other approach\n\nnp.testing.assert_allclose(res1, res2)\n\n%timeit func(theta, pwr, neighbour, coschi, np.ones(size))\n%timeit # other approach\n

\n\n

在我的计算机上，速度已经快得多了：

\n\n

@nb.njit(fastmath=True)\ndef func2(theta, pwr, neighbour, coschi, temp):\n    arccos_coschi = np.arccos(coschi)\n    for k in range(np.argmax(pwr), 5000 * pwr.size): \n        n = k % pwr.size\n        if np.abs(theta[n] - np.pi / 2.) < np.abs(theta_c):\n            adj = neighbour[n, 1]\n        else:\n            adj = neighbour[n, 0]\n        psi_diff = np.abs(arccos_coschi[adj] - arccos_coschi[n])\n        temp5 = temp[adj]**5;\n        e_temp = 1. - np.exp(-temp5 * psi_diff / np.abs(eps))\n        temp[n] = temp[adj] + e_temp / temp5 * (pwr[n] - temp[adj]**4)\n    return temp\n

\n\n

然而，这是有代价的：结果会有所不同！我使用原始版本和提升版本始终得到显着不同的结果fastmath=True。然而结果（几乎）与 相同fastmath=False。似乎fastmath可以实现一些严格的优化，而np.arccos(coschi[adj]) - np.arccos(coschi[n])这些优化在提升时是不可能的np.arccos。在我个人看来，fastmath=True如果您关心确切的结果或者您已经测试过结果的准确性不会受到快速数学的显着影响，我会忽略它！

\n\n

吊装adj

\n\n

接下来要提升的是adj，它的计算次数也比必要的要频繁得多：

\n\n

1.73 s \xc2\xb1 54.1 ms per loop (mean \xc2\xb1 std. dev. of 7 runs, 1 loop each)  # original\n811 ms \xc2\xb1 49.6 ms per loop (mean \xc2\xb1 std. dev. of 7 runs, 1 loop each)  # func2\n

\n\n

这样做的影响不是那么大，但可以衡量：

\n\n

@nb.njit(fastmath=True)\ndef func3(theta, pwr, neighbour, coschi, temp):\n    arccos_coschi = np.arccos(coschi)\n    associated_neighbour = np.empty(neighbour.shape[0], nb.int64)\n    for idx in range(neighbour.shape[0]):\n        if np.abs(theta[idx] - np.pi / 2.) < np.abs(theta_c):\n            associated_neighbour[idx] = neighbour[idx, 1]\n        else:\n            associated_neighbour[idx] = neighbour[idx, 0]\n\n    for k in range(np.argmax(pwr), 5000 * pwr.size): \n        n = k % pwr.size\n        adj = associated_neighbour[n]\n        psi_diff = np.abs(arccos_coschi[adj] - arccos_coschi[n])\n        temp5 = temp[adj]**5;\n        e_temp = 1. - np.exp(-temp5 * psi_diff / np.abs(eps))\n        temp[n] = temp[adj] + e_temp / temp5 * (pwr[n] - temp[adj]**4)\n    return temp\n

\n\n

提升额外的计算似乎对我的计算机的性能没有影响，因此我不将它们包含在此处。所以看起来这就是在不改变算法的情况下你能走多远。

\n\n

重构为更小的函数（+小的额外更改）

\n\n

不过，我建议将其他函数中的提升分开，并使所有变量都成为函数参数，而不是查找全局变量。这可能不会导致加速，但可以使代码更具可读性：

\n\n

1.75 s \xc2\xb1 110 ms per loop (mean \xc2\xb1 std. dev. of 7 runs, 1 loop each)  # original\n761 ms \xc2\xb1 28.1 ms per loop (mean \xc2\xb1 std. dev. of 7 runs, 1 loop each) # func2\n660 ms \xc2\xb1 8.42 ms per loop (mean \xc2\xb1 std. dev. of 7 runs, 1 loop each) # func3\n

\n\n

这里我还做了一些额外的修改：

\n\n

\n
• 预先使用np.tile和 切片而不是 与range一起计算索引%。
\n
• 使用普通 NumPy（numba 之外）来计算np.arccos.
\n

\n\n

最终时间和总结

\n\n

@nb.njit\ndef func4_inner(indices, pwr, associated_neighbour, arccos_coschi, temp, abs_eps):\n    for n in indices:\n        adj = associated_neighbour[n]\n        psi_diff = np.abs(arccos_coschi[adj] - arccos_coschi[n])\n        temp5 = temp[adj]**5;\n        e_temp = 1. - np.exp(-temp5 * psi_diff / abs_eps)\n        temp[n] = temp[adj] + e_temp / temp5 * (pwr[n] - temp[adj]**4)\n    return temp\n\n@nb.njit\ndef get_relevant_neighbor(neighbour, abs_theta_minus_pi_half, abs_theta_c):\n    associated_neighbour = np.empty(neighbour.shape[0], nb.int64)\n    for idx in range(neighbour.shape[0]):\n        if abs_theta_minus_pi_half[idx] < abs_theta_c:\n            associated_neighbour[idx] = neighbour[idx, 1]\n        else:\n            associated_neighbour[idx] = neighbour[idx, 0]\n    return associated_neighbour\n\ndef func4(theta, pwr, neighbour, coschi, temp, theta_c, eps):\n    arccos_coschi = np.arccos(coschi)\n    abs_theta_minus_pi_half = np.abs(theta - (np.pi / 2.))\n    relevant_neighbors = get_relevant_neighbor(neighbour, abs_theta_minus_pi_half, abs(theta_c))\n    argmax_pwr = np.argmax(pwr)\n    indices = np.tile(np.arange(pwr.size), 5000)[argmax_pwr:]\n    return func4_inner(indices, pwr, relevant_neighbors, arccos_coschi, temp, abs(eps))\n

\n\n

所以最终最新的方法大约fastmath比原始方法快 3 倍（没有 ）。如果您确定要使用fastmath，那么只需应用fastmath=True，func4_inner它会更快：

\n\n

1.79 s \xc2\xb1 49.1 ms per loop (mean \xc2\xb1 std. dev. of 7 runs, 1 loop each)  # original\n844 ms \xc2\xb1 41.4 ms per loop (mean \xc2\xb1 std. dev. of 7 runs, 1 loop each)  # func2\n707 ms \xc2\xb1 31.8 ms per loop (mean \xc2\xb1 std. dev. of 7 runs, 1 loop each)  # func3\n550 ms \xc2\xb1 4.88 ms per loop (mean \xc2\xb1 std. dev. of 7 runs, 1 loop each)  # func4\n

\n\n

然而，正如我已经说过的，fastmath如果您想要精确（或至少不是太不精确）的结果，则可能不合适。

\n\n

此外，这里的一些优化在很大程度上取决于可用的硬件和处理器缓存（特别是对于代码的内存带宽有限部分）。您必须检查这些方法在您的计算机上的相对执行情况。

\n