numba 函数中的这个循环可以优化为 运行 更快吗？
Can this loop in a numba function be optimized to run faster?
在我的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秒也意味着为我节省数十小时。
如何进一步优化此代码？
您的示例中似乎处理了很多重复项。
在此版本中，我不会重新计算我们已经看到的 'n' 的任何值。
我不知道这是否适合你的情况，但它为我节省了 ~0.4 秒。
#!/usr/bin/env python
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):
hashtable = {}
for k in range(np.argmax(pwr), 5000*(pwr.size)):
n = k % pwr.size
if not hashtable.get(n, False):
hashtable[n] = 1
#taking into account regions with different super wind direction
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))
retval = temp[adj] + (e_temp)/temp5*(pwr[n] - temp[adj]**4)
temp[n] = retval
return temp
#check time
time1 = time.time()
result = func(theta, pwr, neighbour, coschi, temp)
print("Took: ", time.time()-time1, "
原文：哈希表
2.3726098537445070 : 1.8722639083862305
2.3447792530059814 : 1.9053585529327393
2.3363733291625977 : 1.9104151725769043
2.3447978496551514 : 1.9298338890075684
2.4740016460418700 : 1.9088914394378662
使用 np.ones 的 25M 项的简单循环：
#!/usr/bin/env python
import numba as nb
import numpy as np
import time
temp = np.ones(25000000)
@nb.jit(fastmath=True)
def func(temp):
return [n for n in temp]
time1 = time.time()
result = func(temp)
print("Took: ", time.time()-time1, " seconds for ", len(temp), " items")
花费了：1.2502222061157227 秒用于 25000000 个项目
花费 1.294729232788086 秒处理 25000000 个项目
花费 1.2670648097991943 秒处理 25000000 个项目
花费了：1.2386720180511475 秒 25000000 件商品
花费 1.2517566680908203 秒处理 25000000 件商品
Numba 真的很棒。但是你很绝望，记住你总是可以 write in C (youtube)。在我自己的问题上，我通过逐行转换为 C 获得了比 numba 30% 的性能提升。
如果你想花那么多精力，我建议使用 eigen for vector operations (with vector size know at compile-time) and pybind11，因为它可以在 numpy 和 eigen 之间进行本地转换。当然，请将主循环保持在 Python 中。确保使用适当的编译器标志（如 -O3 -march=native、-mtune=native、-ffast-math）并尝试不同的编译器（对我来说，gcc 输出比之前快 2 倍clang，但同事反映相反。
如果您不懂任何 C++，将自己限制在纯 C 中而不使用库可能更明智（因为它降低了复杂性）。但是您将直接处理 Python 和 numpy C API（不是那么难，但是代码更多，您将了解有关 Python 内部的所有信息）。
让我从一些一般性评论开始：
如果您使用 numba 并且非常关心性能，您应该避免 numba 创建 object-mode 代码的任何可能性。这意味着您应该使用 numba.njit(...) 或 numba.jit(nopython=True, ...) 而不是 numba.jit(...).
这对你的情况没有影响，但它使意图更清晰，并在（快速）nopython 模式不支持某些内容时立即抛出异常。
你应该注意你的时间和方式。第一次调用 numba-jitted 函数（未编译 ahead-of-time）将包含编译成本。所以你需要 运行 它一次，然后再计时以获得准确的时间。为了获得更准确的计时，您应该多次调用该函数。我喜欢 Jupyter Notebooks/Lab 中的 IPythons %timeit 来大致了解性能。
所以我将使用：
res1 = func(theta, pwr, neighbour, coschi, np.ones(size))
res2 = # other approach
np.testing.assert_allclose(res1, res2)
%timeit func(theta, pwr, neighbour, coschi, np.ones(size))
%timeit # other approach
这样我使用断言进行第一次调用（包括编译时间）以确保它确实产生（几乎）相同的输出，然后使用更强大的计时方法（与 time).
吊装np.arccos
现在让我们从一些实际的性能优化开始：一个明显的优化是您可以提升一些 "invariants"，例如 np.arccos(coschi[...]) 的计算频率比 coschi。您对 coschi 中的每个元素进行了大约 5000 次迭代，并且每个循环执行两次 np.arccos！因此，让我们计算一次 coschi 的 arccos 并将其存储在一个中间数组中，以便可以在循环内访问它：
@nb.njit(fastmath=True)
def func2(theta, pwr, neighbour, coschi, temp):
arccos_coschi = np.arccos(coschi)
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(arccos_coschi[adj] - 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
在我的电脑上已经明显快了：
1.73 s ± 54.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each) # original
811 ms ± 49.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each) # func2
然而这是有代价的：结果会不一样！使用 fastmath=True 的原始版本和提升版本，我始终得到明显不同的结果。然而，结果（几乎）等于 fastmath=False。似乎 fastmath 可以对 np.arccos(coschi[adj]) - np.arccos(coschi[n]) 进行一些严格的优化，而当提升 np.arccos 时，这些优化是不可能的。在我个人看来，如果您关心精确的结果或者您已经测试过结果的准确性不受 fastmath 的显着影响，我会忽略 fastmath=True！
吊装adj
接下来要提升的是 adj，它的计算频率也比必要的高得多：
@nb.njit(fastmath=True)
def func3(theta, pwr, neighbour, coschi, temp):
arccos_coschi = np.arccos(coschi)
associated_neighbour = np.empty(neighbour.shape[0], nb.int64)
for idx in range(neighbour.shape[0]):
if np.abs(theta[idx] - np.pi / 2.) < np.abs(theta_c):
associated_neighbour[idx] = neighbour[idx, 1]
else:
associated_neighbour[idx] = neighbour[idx, 0]
for k in range(np.argmax(pwr), 5000 * pwr.size):
n = k % pwr.size
adj = associated_neighbour[n]
psi_diff = np.abs(arccos_coschi[adj] - 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
这个影响不是很大，但是可以衡量：
1.75 s ± 110 ms per loop (mean ± std. dev. of 7 runs, 1 loop each) # original
761 ms ± 28.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each) # func2
660 ms ± 8.42 ms per loop (mean ± std. dev. of 7 runs, 1 loop each) # func3
提升额外的计算似乎对我的计算机性能没有影响，所以我不在这里包括它们。所以这似乎是你可以在不改变算法的情况下达到的程度。
重构为更小的函数（+ 小的附加更改）
但是我建议将提升分离到其他函数中，并使所有变量成为函数参数，而不是查找全局变量。这可能不会导致 speed-ups 但它可以使代码更具可读性：
@nb.njit
def func4_inner(indices, pwr, associated_neighbour, arccos_coschi, temp, abs_eps):
for n in indices:
adj = associated_neighbour[n]
psi_diff = np.abs(arccos_coschi[adj] - arccos_coschi[n])
temp5 = temp[adj]**5;
e_temp = 1. - np.exp(-temp5 * psi_diff / abs_eps)
temp[n] = temp[adj] + e_temp / temp5 * (pwr[n] - temp[adj]**4)
return temp
@nb.njit
def get_relevant_neighbor(neighbour, abs_theta_minus_pi_half, abs_theta_c):
associated_neighbour = np.empty(neighbour.shape[0], nb.int64)
for idx in range(neighbour.shape[0]):
if abs_theta_minus_pi_half[idx] < abs_theta_c:
associated_neighbour[idx] = neighbour[idx, 1]
else:
associated_neighbour[idx] = neighbour[idx, 0]
return associated_neighbour
def func4(theta, pwr, neighbour, coschi, temp, theta_c, eps):
arccos_coschi = np.arccos(coschi)
abs_theta_minus_pi_half = np.abs(theta - (np.pi / 2.))
relevant_neighbors = get_relevant_neighbor(neighbour, abs_theta_minus_pi_half, abs(theta_c))
argmax_pwr = np.argmax(pwr)
indices = np.tile(np.arange(pwr.size), 5000)[argmax_pwr:]
return func4_inner(indices, pwr, relevant_neighbors, arccos_coschi, temp, abs(eps))
这里我也做了一些额外的改动：
- 预先使用
np.tile 和切片而不是 range 方法和 %. 计算了索引
- 使用普通 NumPy（在 numba 之外）计算
np.arccos。
最终时间安排和总结
1.79 s ± 49.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each) # original
844 ms ± 41.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each) # func2
707 ms ± 31.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each) # func3
550 ms ± 4.88 ms per loop (mean ± std. dev. of 7 runs, 1 loop each) # func4
所以最后最新的方法比原来的方法大约快 3 倍（没有 fastmath）。如果您确定要使用 fastmath，那么只需在 func4_inner 上应用 fastmath=True，它会更快：
499 ms ± 4.47 ms per loop (mean ± std. dev. of 7 runs, 1 loop each) # func4 with fastmath on func4_inner
然而，正如我已经说过的，如果您想要准确（或至少不是太不准确）的结果，fastmath 可能不合适。
这里的一些优化在很大程度上取决于可用的硬件和处理器缓存（尤其是 memory-bandwidth-limited 部分代码）。您必须在您的计算机上检查这些方法相对于彼此的执行情况。
在我的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秒也意味着为我节省数十小时。
如何进一步优化此代码？
您的示例中似乎处理了很多重复项。
在此版本中，我不会重新计算我们已经看到的 'n' 的任何值。
我不知道这是否适合你的情况，但它为我节省了 ~0.4 秒。
#!/usr/bin/env python
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):
hashtable = {}
for k in range(np.argmax(pwr), 5000*(pwr.size)):
n = k % pwr.size
if not hashtable.get(n, False):
hashtable[n] = 1
#taking into account regions with different super wind direction
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))
retval = temp[adj] + (e_temp)/temp5*(pwr[n] - temp[adj]**4)
temp[n] = retval
return temp
#check time
time1 = time.time()
result = func(theta, pwr, neighbour, coschi, temp)
print("Took: ", time.time()-time1, "
原文：哈希表
2.3726098537445070 : 1.8722639083862305
2.3447792530059814 : 1.9053585529327393
2.3363733291625977 : 1.9104151725769043
2.3447978496551514 : 1.9298338890075684
2.4740016460418700 : 1.9088914394378662
使用 np.ones 的 25M 项的简单循环：
#!/usr/bin/env python
import numba as nb
import numpy as np
import time
temp = np.ones(25000000)
@nb.jit(fastmath=True)
def func(temp):
return [n for n in temp]
time1 = time.time()
result = func(temp)
print("Took: ", time.time()-time1, " seconds for ", len(temp), " items")
花费了：1.2502222061157227 秒用于 25000000 个项目
花费 1.294729232788086 秒处理 25000000 个项目
花费 1.2670648097991943 秒处理 25000000 个项目
花费了：1.2386720180511475 秒 25000000 件商品
花费 1.2517566680908203 秒处理 25000000 件商品
Numba 真的很棒。但是你很绝望，记住你总是可以 write in C (youtube)。在我自己的问题上，我通过逐行转换为 C 获得了比 numba 30% 的性能提升。
如果你想花那么多精力，我建议使用 eigen for vector operations (with vector size know at compile-time) and pybind11，因为它可以在 numpy 和 eigen 之间进行本地转换。当然，请将主循环保持在 Python 中。确保使用适当的编译器标志（如 -O3 -march=native、-mtune=native、-ffast-math）并尝试不同的编译器（对我来说，gcc 输出比之前快 2 倍clang，但同事反映相反。
如果您不懂任何 C++，将自己限制在纯 C 中而不使用库可能更明智（因为它降低了复杂性）。但是您将直接处理 Python 和 numpy C API（不是那么难，但是代码更多，您将了解有关 Python 内部的所有信息）。
让我从一些一般性评论开始：
如果您使用 numba 并且非常关心性能，您应该避免 numba 创建 object-mode 代码的任何可能性。这意味着您应该使用
numba.njit(...)或numba.jit(nopython=True, ...)而不是numba.jit(...).这对你的情况没有影响，但它使意图更清晰，并在（快速）nopython 模式不支持某些内容时立即抛出异常。
你应该注意你的时间和方式。第一次调用 numba-jitted 函数（未编译 ahead-of-time）将包含编译成本。所以你需要 运行 它一次，然后再计时以获得准确的时间。为了获得更准确的计时，您应该多次调用该函数。我喜欢 Jupyter Notebooks/Lab 中的 IPythons
%timeit来大致了解性能。所以我将使用：
res1 = func(theta, pwr, neighbour, coschi, np.ones(size)) res2 = # other approach np.testing.assert_allclose(res1, res2) %timeit func(theta, pwr, neighbour, coschi, np.ones(size)) %timeit # other approach这样我使用断言进行第一次调用（包括编译时间）以确保它确实产生（几乎）相同的输出，然后使用更强大的计时方法（与
time).
吊装np.arccos
现在让我们从一些实际的性能优化开始：一个明显的优化是您可以提升一些 "invariants"，例如 np.arccos(coschi[...]) 的计算频率比 coschi。您对 coschi 中的每个元素进行了大约 5000 次迭代，并且每个循环执行两次 np.arccos！因此，让我们计算一次 coschi 的 arccos 并将其存储在一个中间数组中，以便可以在循环内访问它：
@nb.njit(fastmath=True)
def func2(theta, pwr, neighbour, coschi, temp):
arccos_coschi = np.arccos(coschi)
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(arccos_coschi[adj] - 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
在我的电脑上已经明显快了：
1.73 s ± 54.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each) # original
811 ms ± 49.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each) # func2
然而这是有代价的：结果会不一样！使用 fastmath=True 的原始版本和提升版本，我始终得到明显不同的结果。然而，结果（几乎）等于 fastmath=False。似乎 fastmath 可以对 np.arccos(coschi[adj]) - np.arccos(coschi[n]) 进行一些严格的优化，而当提升 np.arccos 时，这些优化是不可能的。在我个人看来，如果您关心精确的结果或者您已经测试过结果的准确性不受 fastmath 的显着影响，我会忽略 fastmath=True！
吊装adj
接下来要提升的是 adj，它的计算频率也比必要的高得多：
@nb.njit(fastmath=True)
def func3(theta, pwr, neighbour, coschi, temp):
arccos_coschi = np.arccos(coschi)
associated_neighbour = np.empty(neighbour.shape[0], nb.int64)
for idx in range(neighbour.shape[0]):
if np.abs(theta[idx] - np.pi / 2.) < np.abs(theta_c):
associated_neighbour[idx] = neighbour[idx, 1]
else:
associated_neighbour[idx] = neighbour[idx, 0]
for k in range(np.argmax(pwr), 5000 * pwr.size):
n = k % pwr.size
adj = associated_neighbour[n]
psi_diff = np.abs(arccos_coschi[adj] - 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
这个影响不是很大，但是可以衡量：
1.75 s ± 110 ms per loop (mean ± std. dev. of 7 runs, 1 loop each) # original
761 ms ± 28.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each) # func2
660 ms ± 8.42 ms per loop (mean ± std. dev. of 7 runs, 1 loop each) # func3
提升额外的计算似乎对我的计算机性能没有影响，所以我不在这里包括它们。所以这似乎是你可以在不改变算法的情况下达到的程度。
重构为更小的函数（+ 小的附加更改）
但是我建议将提升分离到其他函数中，并使所有变量成为函数参数，而不是查找全局变量。这可能不会导致 speed-ups 但它可以使代码更具可读性：
@nb.njit
def func4_inner(indices, pwr, associated_neighbour, arccos_coschi, temp, abs_eps):
for n in indices:
adj = associated_neighbour[n]
psi_diff = np.abs(arccos_coschi[adj] - arccos_coschi[n])
temp5 = temp[adj]**5;
e_temp = 1. - np.exp(-temp5 * psi_diff / abs_eps)
temp[n] = temp[adj] + e_temp / temp5 * (pwr[n] - temp[adj]**4)
return temp
@nb.njit
def get_relevant_neighbor(neighbour, abs_theta_minus_pi_half, abs_theta_c):
associated_neighbour = np.empty(neighbour.shape[0], nb.int64)
for idx in range(neighbour.shape[0]):
if abs_theta_minus_pi_half[idx] < abs_theta_c:
associated_neighbour[idx] = neighbour[idx, 1]
else:
associated_neighbour[idx] = neighbour[idx, 0]
return associated_neighbour
def func4(theta, pwr, neighbour, coschi, temp, theta_c, eps):
arccos_coschi = np.arccos(coschi)
abs_theta_minus_pi_half = np.abs(theta - (np.pi / 2.))
relevant_neighbors = get_relevant_neighbor(neighbour, abs_theta_minus_pi_half, abs(theta_c))
argmax_pwr = np.argmax(pwr)
indices = np.tile(np.arange(pwr.size), 5000)[argmax_pwr:]
return func4_inner(indices, pwr, relevant_neighbors, arccos_coschi, temp, abs(eps))
这里我也做了一些额外的改动：
- 预先使用
np.tile和切片而不是range方法和%. 计算了索引
- 使用普通 NumPy（在 numba 之外）计算
np.arccos。
最终时间安排和总结
1.79 s ± 49.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each) # original
844 ms ± 41.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each) # func2
707 ms ± 31.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each) # func3
550 ms ± 4.88 ms per loop (mean ± std. dev. of 7 runs, 1 loop each) # func4
所以最后最新的方法比原来的方法大约快 3 倍（没有 fastmath）。如果您确定要使用 fastmath，那么只需在 func4_inner 上应用 fastmath=True，它会更快：
499 ms ± 4.47 ms per loop (mean ± std. dev. of 7 runs, 1 loop each) # func4 with fastmath on func4_inner
然而，正如我已经说过的，如果您想要准确（或至少不是太不准确）的结果，fastmath 可能不合适。
这里的一些优化在很大程度上取决于可用的硬件和处理器缓存（尤其是 memory-bandwidth-limited 部分代码）。您必须在您的计算机上检查这些方法相对于彼此的执行情况。