Question

我正在实现一种具有特定结构的顺序算法（卡尔曼滤波器），其中许多内部循环可以并行完成。我需要从该功能中获得尽可能多的性能。目前，它在具有代表性数据输入（n，p = 12，d = 3，T = 3000）的计算机上运行时间约为600ms

我已将@numba.jit与nopython=True, parallel=True一起使用，并为numba.prange注释了我的范围。但是，即使输入的数据非常多（n> 5000），也显然不会发生并行性（基于仅查看top的内核）。

这里有很多代码，我只显示主要块。 Numba无法在prange下并行化数组操作是有原因的吗？我还检查了numba.config.NUMBA_NUM_THREADS（它是8）并使用了不同的numba.config.THREADING_LAYER（当前是'tbb'）。我也尝试过使用openpys和numpy + scipy的MKL版本，MKL版本似乎稍慢一些，并且仍然没有并行化。

注释为：

@numba.jit(nopython=True, cache=False, parallel=True,
           fastmath=True, nogil=True)

函数的主要部分：

P = np.empty((T + 1, n, p, d, d))
m = np.empty((T + 1, n, p, d))

P[0] = P0
m[0] = m0

phi = 0.0
Xt = np.empty((n, p)

for t in range(1, T + 1):
    sum_P00 = 0.0
    v = y[t - 1]

    # Purely for convenience, little performance impact
    for tau in range(1, p + 1):
        Xt[:, tau - 1] = X[p + t - 1 - tau]

    # Predict
    for i in numba.prange(n):
        for tau in range(p):
            # Prediction step
            m[t, i, tau] = Phi[i, tau] @ m[t - 1, i, tau]
            P[t, i, tau] = Phi[i, tau] @ P[t - 1, i, tau] @ Phi[i, tau].T

    # Auxiliary gain variables
    for i in numba.prange(n):
        for tau in range(p):
            v = v - Xt[i, tau] * m[t, i, tau, 0]
            sum_P00 = sum_P00 + P[t, i, tau, 0, 0]

    # Energy function update
    s = np.linalg.norm(Xt)**2 * sum_P00 + sv2
    phi += np.pi * s + 0.5 * v**2 / s

    # Update
    for i in numba.prange(n):
        for tau in range(p):
            k = Xt[i, tau] * P[t, i, tau, :, 0]  # Gain
            m[t, i, tau] = m[t, i, tau] + (v / s) * k
            P[t, i, tau] = P[t, i, tau] + (k / s) @ k.T

Answer 1

在Ipython中以交互方式运行似乎只是一个问题。如预期的那样，从控制台运行测试脚本会导致并行执行。

使用Numba进行并行循环-不使用prange进行并行化

1 个答案: