加速卷积函数的 Numba 优化实战

本文旨在指导如何使用 Numba 优化卷积函数，通过将 NumPy 代码替换为显式循环，并利用 Numba 的并行化功能，显著提升代码执行效率。我们将深入探讨优化策略，并提供优化后的代码示例，最终实现比原始 NumPy 代码快数倍的加速效果。

优化思路：避免临时数组和利用显式循环

原始代码中使用了大量的 NumPy 操作，这在 Numba 中可能会导致性能瓶颈。Numba 更擅长处理显式循环，因为它能够更好地进行编译优化。此外，频繁创建临时数组会带来额外的开销，影响性能。因此，优化的关键在于：

1. 使用显式循环代替 NumPy 操作：将 np.dot 等 NumPy 函数替换为等效的循环实现。
2. 减少临时数组的创建：尽可能在循环内部进行计算，避免生成不必要的临时数组。
3. 利用 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

代码解释：

腾讯智影-AI数字人

基于AI数字人能力，实现7*24小时AI数字人直播带货，低成本实现直播业务快速增增，全天智能在线直播

73

查看详情

• @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 加速等技术，进一步提升代码的执行效率。

以上就是加速卷积函数的 Numba 优化实战的详细内容，更多请关注php中文网其它相关文章！