循环中的向量化比 numba jited 函数中的嵌套循环慢

Question

所以我正在尝试在numba（我目前使用numba 0.45.1）中结合@njit 支持的矢量化和for循环来提高性能。令人失望的是，我发现它实际上比我的代码中的纯嵌套循环实现要慢。

这是我的代码：

import numpy as np
from numba import njit

@njit
def func3(arr_in, win_arr):
    n = arr_in.shape[0]
    win_len = len(win_arr)

    result = np.full((n, win_len), np.nan)

    alpha_arr = 2 / (win_arr + 1)

    e = np.full(win_len, arr_in[0])
    w = np.ones(win_len)

    two_index = np.nonzero(win_arr <= 2)[0][-1]+1
    result[0, :two_index] = arr_in[0]

    for i in range(1, n):
        w = w + (1-alpha_arr)**i
        e = e*(1-alpha_arr) + arr_in[i]
        result[i,:] = e /w

    return result

@njit
def func4(arr_in, win_arr):
    n = arr_in.shape[0]
    win_len = len(win_arr)

    result = np.full((n, win_len), np.nan)

    alpha_arr = 2 / (win_arr + 1)

    e = np.full(win_len, arr_in[0])
    w = np.ones(win_len)

    two_index = np.nonzero(win_arr <= 2)[0][-1]+1
    result[0, :two_index] = arr_in[0]

    for i in range(1, n):
        for col in range(len(win_arr)):
            w[col] = w[col] + (1-alpha_arr[col])**i
            e[col] = e[col]*(1-alpha_arr[col]) + arr_in[i]
            result[i,col] = e[col] /w[col]

    return result

if __name__ == '__main__':
    np.random.seed(0)
    data_size = 200000
    winarr_size = 1000

    data = np.random.uniform(0,1000, size = data_size)+29000
    win_array = np.arange(1, winarr_size+1)

    abc_test3= func3(data, win_array)
    abc_test4= func4(data, win_array)

    print(np.allclose(abc_test3, abc_test4, equal_nan = True))

我使用以下配置对这两个函数进行了基准测试：

(data_size,winarr_size) = (200000,100), (200000,200),(200000,1000), (200000,2000), (20000,10000), (2000,100000).

并发现纯嵌套 for-loop 实现 (func4) 始终比 for-loop 与矢量化混合实现 (func3) 更快（大约快 2-5%）。

---

我的问题如下：

1) 进一步提高代码速度需要做哪些改动？

2）为什么函数的向量化版本的计算时间随着win_arr的大小线性增长？我认为矢量化应该使得无论矢量有多大/多小，运算速度都是恒定的，但显然这在这种情况下并不成立。

3) 是否存在一般情况下矢量化运算的计算时间仍会随输入大小线性增长？

Answer 1

对于复杂的数学运算，例如（pow、除法、...），请务必三思而后行。如果您可以用简单的运算（例如乘法、加法和减法）来替换它们，那总是值得一试的。

请注意，将 alpha 与自身重复相乘仅在代数上与直接用幂计算相同。由于这是数值数学，因此结果可能会有所不同。

还要避免不必要的临时数组。

第一次尝试

@nb.njit(error_model="numpy",parallel=True)
def func5(arr_in, win_arr):
    #filling the whole array with NaNs isn't necessary
    result = np.empty((win_arr.shape[0],arr_in.shape[0]))
    for col in range(win_arr.shape[0]):
        result[col,0]=np.nan

    two_index = np.nonzero(win_arr <= 2)[0][-1]+1
    result[:two_index,0] = arr_in[0]

    for col in nb.prange(win_arr.shape[0]):
        alpha=1.-(2./ (win_arr[col] + 1.))
        alpha_exp=alpha

        w=1.
        e=arr_in[0]

        for i in range(1, arr_in.shape[0]):
            w+= alpha_exp
            e = e*alpha + arr_in[i]
            result[col,i] = e/w
            alpha_exp*=alpha

    return result.T

第二次尝试（避免下溢）

@nb.njit(error_model="numpy",parallel=True)
def func7(arr_in, win_arr):
    #filling the whole array with NaNs isn't necessary
    result = np.empty((win_arr.shape[0],arr_in.shape[0]))
    for col in range(win_arr.shape[0]):
        result[col,0]=np.nan

    two_index = np.nonzero(win_arr <= 2)[0][-1]+1
    result[:two_index,0] = arr_in[0]

    for col in nb.prange(win_arr.shape[0]):
        alpha=1.-(2./ (win_arr[col] + 1.))
        alpha_exp=alpha

        w=1.
        e=arr_in[0]

        for i in range(1, arr_in.shape[0]):
            w+= alpha_exp
            e = e*alpha + arr_in[i]
            result[col,i] = e/w

          if np.abs(alpha_exp)>=1e-308:
              alpha_exp*=alpha
          else:
              alpha_exp=0.

    return result.T

时间

%timeit abc_test3= func3(data, win_array)
7.17 s ± 45.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit abc_test4= func4(data, win_array)
7.13 s ± 13.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
#from MSeifert answer (parallelized)
%timeit abc_test6= func6(data, win_array)
3.42 s ± 153 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit abc_test5= func5(data, win_array)
1.22 s ± 22.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit abc_test7= func7(data, win_array)
238 ms ± 5.55 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Answer 2

您似乎误解了“矢量化”的含义。向量化意味着您编写的代码对数组进行操作，就好像它们是标量一样——但这就是代码的样子，与性能无关。

在 Python/NumPy 世界中，向量化还意味着与循环代码相比，向量化操作中的循环开销（通常）要小得多。但是矢量化代码仍然需要执行循环（即使它隐藏在库中）！

此外，如果您使用 numba 编写循环，numba 将编译它并创建执行速度（通常）与矢量化 NumPy 代码一样快的代码。这意味着在 numba 函数内部，矢量化代码和非矢量化代码之间没有显着的性能差异。

所以这应该回答你的问题：

2）为什么函数的向量化版本的计算时间随着win_arr的大小线性增长？我认为矢量化应该使得无论矢量有多大/多小，操作速度都是恒定的，但显然这在这种情况下并不成立。

它线性增长，因为它仍然需要迭代。在矢量化代码中，循环只是隐藏在库例程中。

3) 是否存在一般情况下矢量化运算的计算时间仍会随输入大小线性增长？

没有。

---

您还询问了可以做些什么来加快速度。

cmets 已经提到你可以并行化它：

import numpy as np
import numba as nb

@nb.njit(parallel=True)
def func6(arr_in, win_arr):
    n = arr_in.shape[0]
    win_len = len(win_arr)

    result = np.full((n, win_len), np.nan)

    alpha_arr = 2 / (win_arr + 1)

    e = np.full(win_len, arr_in[0])
    w = np.ones(win_len)

    two_index = np.nonzero(win_arr <= 2)[0][-1]+1
    result[0, :two_index] = arr_in[0]

    for i in range(1, n):
        for col in nb.prange(len(win_arr)):
            w[col] = w[col] + (1-alpha_arr[col])**i
            e[col] = e[col] * (1-alpha_arr[col]) + arr_in[i]
            result[i,col] = e[col] /w[col]

    return result

这使我的机器上的代码速度更快（4 核）。

但是还有一个问题是您的算法可能在数值上不稳定。当你将(1-alpha_arr[col])**i 提升到十万次方时，它会在某个时候下溢：

>>> alpha = 0.01
>>> for i in [1, 10, 100, 1_000, 10_000, 50_000, 100_000, 200_000]:
...     print((1-alpha)**i)
0.99
0.9043820750088044
0.3660323412732292
4.317124741065786e-05
2.2487748498162805e-44
5.750821364590612e-219
0.0  # <-- underflow
0.0