我试图用一个简单的例子来说明为我自己和一些同事使用numba.prange
的好处,但我无法获得适当的加速.我编写了一个简单的一维扩散解算器,它基本上在一个长数组上循环,组合元素i+1
、i
和i-1
,并将结果写入第二个数组的元素i
.这应该是一个非常完美的并行for循环用例,类似于Fortran或C中的OpenMP.
下面是我的完整示例:
import numpy as np
from numba import jit, prange
@jit(nopython=True, parallel=True)
def diffusion(Nt):
alpha = 0.49
x = np.linspace(0, 1, 10000000)
# Initial condition
C = 1/(0.25*np.sqrt(2*np.pi)) * np.exp(-0.5*((x-0.5)/0.25)**2)
# Temporary work array
C_ = np.zeros_like(C)
# Loop over time (normal for-loop)
for j in range(Nt):
# Loop over array elements (space, parallel for-loop)
for i in prange(1, len(C)-1):
C_[i] = C[i] + alpha*(C[i+1] - 2*C[i] + C[i-1])
C[:] = C_
return C
# Run once to just-in-time compile
C = diffusion(1)
# Check timing
%timeit C = diffusion(100)
当使用parallel=False
时,这大约需要2秒,使用parallel=True
时大约需要1.5秒.我在一台有4个物理内核的MacBook Pro上运行,Activity Monitor报告在有无并行化的情况下CPU使用率分别为parallel=False
%和700%左右.
我本以为会有更接近因子4的加速.我做错什么了吗?