我正在try 使用numba并行化一个简单的python程序.它由两个功能组成.第一种是简单幂法
@numba.jit(nopython=True)
def power_method(A, v):
u = v.copy()
for i in range(3 * 10**3):
u = A @ u
u /= np.linalg.norm(u)
return u
第二个函数在网格上与向量v迭代，并为各种向量v运行power_method函数.
@numba.jit(nopython=True, parallel=True)
def iterate_grid(A, scale, sz):
assert A.shape[0] == A.shape[1] == 3
n = A.shape[0]
results = np.empty((sz**3, n))
tmp = np.linspace(-scale, scale, sz)
for i1 in range(sz):
v = np.empty(n, dtype=np.float64)
v1 = tmp[i1]
for i2, v2 in enumerate(tmp):
for i3, v3 in enumerate(tmp):
v[0] = v1
v[1] = v2
v[2] = v3
u = power_method(A, v)
idx = i1 * sz**2 + i2 * sz + i3
results[idx] = u.copy()
return results
然后我用
n = 3
A = np.random.randn(n, n)
iterate_grid(A, 5.0, 20)
所有迭代都是独立的.此外，理想情况下，计算会落在缓存中，因此我希望将第一个循环与prange并行将得到近似的线性加速度.
然而，顺序代码的壁时间为6.07秒，而并行代码的壁时间为6.07秒
@numba.jit(nopython=True, parallel=True)
def iterate_grid(A, scale, sz):
assert A.shape[0] == A.shape[1] == 3
n = A.shape[0]
results = np.empty((sz**3, n))
tmp = np.linspace(-scale, scale, sz)
for i1 in numba.prange(sz):
v = np.empty(n, dtype=np.float64)
v1 = tmp[i1]
for i2, v2 in enumerate(tmp):
for i3, v3 in enumerate(tmp):
v[0] = v1
v[1] = v2
v[2] = v3
u = power_method(A, v)
idx = i1 * sz**2 + i2 * sz + i3
results[idx] = u.copy()
return results
挂机时间为7.79秒.也就是说，在这种情况下，并行化会减慢代码的速度.
此外，正如我从iterate_grid.parallel_diagnostics(level=4)中看到的，numba熔断了
并行化power方法并不是我真正想要解决的任务，它只是一个小例子，numba的行为并没有达到我的预期.你能给我解释一下numba的这种行为，并建议我如何在网格上并行化迭代以获得线性加速度吗？
提前感谢您的帮助！