本文旨在指导如何使用 Numba 优化卷积函数，通过将 NumPy 代码替换为显式循环，并利用 Numba 的并行化功能，显著提升代码执行效率。我们将深入探讨优化策略，并提供优化后的代码示例，最终实现比原始 NumPy 代码快数倍的加速效果。
优化思路：避免临时数组和利用显式循环
原始代码中使用了大量的 NumPy 操作，这在 Numba 中可能会导致性能瓶颈。Numba 更擅长处理显式循环，因为它能够更好地进行编译优化。此外，频繁创建临时数组会带来额外的开销，影响性能。因此，优化的关键在于：
- 使用显式循环代替 NumPy 操作：将 np.dot 等 NumPy 函数替换为等效的循环实现。
- 减少临时数组的创建：尽可能在循环内部进行计算，避免生成不必要的临时数组。
- 利用 Numba 的并行化功能：通过 nb.prange 并行化外层循环，充分利用多核 CPU 的计算能力。
优化后的代码示例
以下是优化后的代码示例，该代码使用显式循环代替 NumPy 操作，并利用 nb.prange 并行化外层循环：
import numpy as np
import numba as nb
@nb.jit(nopython=True, parallel=True)
def numba_convolve_faster(wvl_sensor, fwhm_sensor, wvl_lut, rad_lut):
num_chans, num_col = wvl_sensor.shape
num_bins = wvl_lut.shape[0]
num_rad = rad_lut.shape[0]
original_res = np.empty((num_col, num_rad, num_chans), dtype=np.float64)
sigma = fwhm_sensor / (2.0 * np.sqrt(2.0 * np.log(2.0)))
var = sigma ** 2
denom = (2 * np.pi * var) ** 0.5
inv_denom = 1.0 / denom
factor = -1 / (2*var)
for x in nb.prange(wvl_sensor.shape[1]):
wvl_sensor_col = wvl_sensor[:, x].copy()
response = np.empty(num_bins)
for j in range(num_chans):
response_sum = 0.0
for i in range(num_bins):
diff = wvl_lut[i] - wvl_sensor_col[j]
response[i] = np.exp(diff * diff * factor[j]) * inv_denom[j]
response_sum += response[i]
inv_response_sum = 1.0 / response_sum
for i in range(num_bins):
response[i] *= inv_response_sum
for k in range(num_rad):
s = 0.0
for i in range(num_bins):
s += rad_lut[k, i] * response[i]
original_res[x, k, j] = s
return original_res代码解释：
- @nb.jit(nopython=True, parallel=True): 使用 Numba 的 jit 装饰器，并开启 nopython 模式和并行化。nopython=True 确保代码完全在 Numba 中编译执行，避免回退到 Python 解释器，从而获得最佳性能。parallel=True 允许 Numba 并行化循环。
- nb.prange(wvl_sensor.shape[1]): 使用 nb.prange 代替 range，允许 Numba 并行化外层循环。
- 显式循环：使用显式循环代替 np.dot 等 NumPy 操作，减少临时数组的创建。
- 预先计算：将一些不变的量（如 sigma, var, denom, inv_denom, factor）在循环外部计算好，避免重复计算。
- .copy(): 在循环内部使用 wvl_sensor[:, x].copy() 创建 wvl_sensor_col 的副本，避免潜在的内存访问问题。
性能提升
经过优化，代码的执行效率得到了显著提升。在 i5-9600KF CPU 上，优化后的代码比原始 NumPy 代码快 5.74 倍。
注意事项
- 数据类型：确保输入数据的类型与 Numba 兼容。建议使用 NumPy 数组，并指定明确的数据类型（如 np.float64）。
- Numba 编译：Numba 首次编译代码时会花费一些时间。后续调用将使用缓存的编译版本，速度会更快。
- 并行化调试：并行化可能会引入一些难以调试的问题。建议先在单线程模式下验证代码的正确性，然后再开启并行化。
- NumPy 函数的限制：并非所有的 NumPy 函数都可以在 Numba 的 nopython 模式下使用。如果遇到不支持的函数，需要使用等效的循环实现。
总结
通过将 NumPy 代码替换为显式循环，并利用 Numba 的并行化功能，可以显著提升卷积函数的执行效率。在实际应用中，可以根据具体情况调整优化策略，以获得最佳性能。此外，还可以考虑使用单精度浮点数、GPU 加速等技术，进一步提升代码的执行效率。
