本文旨在指导读者如何使用 Numba 优化卷积函数，通过避免创建临时数组、采用显式循环以及利用 Numba 的并行计算能力，显著提升代码执行效率。我们将对比原始 NumPy 实现和优化后的 Numba 实现，并深入探讨优化策略背后的原理，最终实现高达 5.74 倍的性能提升。
问题分析与优化思路
原始的卷积函数实现依赖于 NumPy 的广播机制和向量化操作，虽然代码简洁，但在大规模数据处理时会产生大量的临时数组，导致性能瓶颈。Numba 可以将 Python 代码编译成机器码，从而加速数值计算。然而，直接使用 Numba 编译原始代码并不能达到理想的效果，因为 NumPy 的一些高级特性在 Numba 并行编译中可能会出现问题。
因此，优化思路主要集中在以下几个方面：
- 避免创建临时数组： 使用显式循环代替 NumPy 的广播机制，减少内存分配和数据拷贝。
- 利用 Numba 的并行计算能力： 使用 nb.prange 并行化外层循环，充分利用多核 CPU 的计算资源。
- 避免使用 BLAS 库： np.dot 函数底层调用 BLAS 库，在高并发场景下可能造成性能瓶颈，使用显式循环代替。
优化后的 Numba 实现
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 模式强制 Numba 将代码编译成纯机器码，避免回退到 Python 解释器，从而获得最佳性能。parallel=True 启用自动并行化。
- 预先计算 sigma，var，denom，inv_denom 和 factor，避免在循环中重复计算。
- 使用显式循环代替 NumPy 的广播机制和 np.dot 函数，减少临时数组的创建。
- wvl_sensor_col = wvl_sensor[:, x].copy()：在并行循环中，确保每个线程都有自己的数据副本，避免数据竞争。
性能测试
在 i5-9600KF CPU (6 cores) 上的测试结果如下：
| 实现方式 | 耗时 |
|---|---|
| NumPy 代码 | 4 分 37 秒 |
| 原始 Numba 代码 | 4 分 24 秒 |
| 优化后的 Numba 代码 | 46 秒 |
优化后的 Numba 实现比原始 NumPy 实现快 5.74 倍。
注意事项与总结
- 数据类型： Numba 对数据类型非常敏感，尽量使用相同的数据类型，避免类型转换带来的性能损失。
- Numba 的局限性： Numba 并非万能，对于复杂的 Python 代码，Numba 可能无法进行有效优化。
- GPU 加速： 对于计算密集型任务，可以考虑使用 GPU 加速，但需要注意数据传输的开销。
- 单精度： 如果精度要求不高，可以使用单精度浮点数，可以显著提升计算速度。
通过本次优化，我们学习了如何使用 Numba 加速卷积函数，掌握了避免创建临时数组、使用显式循环以及利用 Numba 的并行计算能力等优化技巧。这些技巧不仅适用于卷积函数，也适用于其他数值计算任务。希望本文能够帮助读者更好地利用 Numba 提升代码性能。
