我正在实现一种具有特定结构的顺序算法(卡尔曼滤波器),其中可以并行完成大量内部循环。我需要从这个函数中获得尽可能多的性能。目前,它在我的机器上运行大约 600 毫秒,具有代表性数据输入(n,p = 12,d = 3,T = 3000)
我已将 @numba.jit
与 nopython=True, parallel=True
一起使用,并使用 numba.prange
注释了我的范围。然而,即使有非常大的数据输入 (n > 5000),显然也不会发生并行性(仅基于使用 top
查看内核)。
这里有相当多的代码,我只显示主要 block 。 Numba 无法在 prange
下并行化数组操作是否有原因?我还检查了 numba.config.NUMBA_NUM_THREADS
(它是 8)并使用不同的 numba.config.THREADING_LAYER
(目前是 'tbb'
)。我也尝试过 openblas 和 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
最佳答案
这似乎只是在 Ipython 中交互运行的问题。正如预期的那样,从控制台运行测试脚本会导致并行执行。
关于python - 与 Numba 并行循环——与 prange 不并行,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/54752254/