Python 如何使numba@jit使用所有cpu核（并行化numba@jit）

我使用numbas

@jit

decorator在python中添加两个numpy数组。如果我使用

@jit

与

python

相比，性能是如此之高

然而，即使我传入了

@numba.jit（nopython=True，parallel=True，nogil=True），它也没有利用所有的CPU核

是否有任何方法可以使用所有具有numba
@jit
的CPU内核

这是我的密码：

import time                                                
import numpy as np                                         
import numba                                               

SIZE = 2147483648 * 6                                      

a = np.full(SIZE, 1, dtype = np.int32)                     

b = np.full(SIZE, 1, dtype = np.int32)                     

c = np.ndarray(SIZE, dtype = np.int32)                     

@numba.jit(nopython = True, parallel = True, nogil = True) 
def add(a, b, c):                                          
    for i in range(SIZE):                                  
        c[i] = a[i] + b[i]                                 

start = time.time()                                        
add(a, b, c)                                               
end = time.time()                                          

print(end - start)                                        

您可以将
parallel=True
传递给任何numba jitted函数，但这并不意味着它总是使用所有内核。您必须理解，numba使用一些启发式方法使代码并行执行，有时这些启发式方法在代码中根本找不到任何可并行的东西。目前存在一个问题，因此如果无法使其“并行”，它会发出警告。因此，它更像是一个“请尽可能使其并行执行”参数，而不是“强制并行执行”

但是，如果您确实知道可以并行化代码，那么您总是可以手动使用线程或进程。只是调整一下：

我突出显示了我更改的部分，其他所有内容都是从示例中逐字复制的。这利用了我机器上的所有内核（4核机器因此有4个线程），但没有显示出显著的加速：

numpy (1 thread)       539 ms
numba (1 thread)       536 ms
numba (4 threads)      442 ms

在这种情况下，多线程处理缺乏（很多）加速是因为加法是一种带宽有限的操作。这意味着从数组中加载元素并将结果放入结果数组所花费的时间要比实际相加所花费的时间多得多

在这些情况下，您甚至可以看到由于并行执行而导致的速度减慢

只有当函数更加复杂，并且与加载和存储数组元素相比，实际操作需要花费大量时间时，您才会看到并行执行的巨大改进。numba文档中的示例如下所示：

def func_np(a, b):
    """
    Control function using Numpy.
    """
    return np.exp(2.1 * a + 3.2 * b)

@jit('void(double[:], double[:], double[:])', nopython=True, nogil=True)
def inner_func_nb(result, a, b):
    """
    Function under test.
    """
    for i in range(len(result)):
        result[i] = math.exp(2.1 * a[i] + 3.2 * b[i])

这实际上（几乎）与线程数成正比，因为两次乘法、一次加法和一次调用
math.exp
要比加载和存储结果慢得多：

func_nb = make_singlethread(inner_func_nb)
func_nb_mt2 = make_multithread(inner_func_nb, 2)
func_nb_mt3 = make_multithread(inner_func_nb, 3)
func_nb_mt4 = make_multithread(inner_func_nb, 4)

a = np.random.rand(size)
b = np.random.rand(size)

correct = timefunc(None, "numpy (1 thread)", func_np, a, b)
timefunc(correct, "numba (1 thread)", func_nb, a, b)
timefunc(correct, "numba (2 threads)", func_nb_mt2, a, b)
timefunc(correct, "numba (3 threads)", func_nb_mt3, a, b)
timefunc(correct, "numba (4 threads)", func_nb_mt4, a, b)

结果:

numpy (1 thread)      3422 ms
numba (1 thread)      2959 ms
numba (2 threads)     1555 ms
numba (3 threads)     1080 ms
numba (4 threads)      797 ms

为了完整性起见，2018年（numba v 0.39）您可以

from numba import prange

并将原始函数定义中的
range
替换为
prange
，就是这样

这立即使CPU利用率达到100%，在我的例子中，运行时间从2.9秒提高到1.7秒（对于大小=2147483648*1，在16核32线程的机器上）

更复杂的内核通常可以通过传入
fastmath=True
来加快速度。您的示例没有使用
prange
——我找到了您的问题，因为我想知道为什么
prange
不能并行运行，但我在
@jit
中缺少了
parallel=True
。。。
from numba import prange