我正在使用numbas@jit装饰器在python中添加两个numpy数组。与python相比，如果我使用@jit，性能是如此之高。
但是，即使我传入@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抖动函数，但这并不意味着它总是利用所有内核。你必须明白，numba 使用一些启发式方法使代码并行执行，有时这些启发式方法在代码中根本找不到任何可以并行化的东西。当前有一个拉取请求，以便在无法使其"并行"时发出警告。所以它更像是一个"如果可能的话，请让它并行执行"参数，而不是"强制并行执行"。
但是，如果您真的知道可以并行化代码，则始终可以手动使用线程或进程。只是改编了 numba 文档中使用多线程的示例：
#!/usr/bin/env python
from __future__ import print_function, division, absolute_import
import math
import threading
from timeit import repeat
import numpy as np
from numba import jit
nthreads = 4
size = 10**7 # CHANGED
# CHANGED
def func_np(a, b):
"""
Control function using Numpy.
"""
return a + b
# CHANGED
@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] = a[i] + b[i]
def timefunc(correct, s, func, *args, **kwargs):
"""
Benchmark *func* and print out its runtime.
"""
print(s.ljust(20), end=" ")
# Make sure the function is compiled before we start the benchmark
res = func(*args, **kwargs)
if correct is not None:
assert np.allclose(res, correct), (res, correct)
# time it
print('{:>5.0f} ms'.format(min(repeat(lambda: func(*args, **kwargs),
number=5, repeat=2)) * 1000))
return res
def make_singlethread(inner_func):
"""
Run the given function inside a single thread.
"""
def func(*args):
length = len(args[0])
result = np.empty(length, dtype=np.float64)
inner_func(result, *args)
return result
return func
def make_multithread(inner_func, numthreads):
"""
Run the given function inside *numthreads* threads, splitting its
arguments into equal-sized chunks.
"""
def func_mt(*args):
length = len(args[0])
result = np.empty(length, dtype=np.float64)
args = (result,) + args
chunklen = (length + numthreads - 1) // numthreads
# Create argument tuples for each input chunk
chunks = [[arg[i * chunklen:(i + 1) * chunklen] for arg in args]
for i in range(numthreads)]
# Spawn one thread per chunk
threads = [threading.Thread(target=inner_func, args=chunk)
for chunk in chunks]
for thread in threads:
thread.start()
for thread in threads:
thread.join()
return result
return func_mt
func_nb = make_singlethread(inner_func_nb)
func_nb_mt = make_multithread(inner_func_nb, nthreads)
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 (%d threads)" % nthreads, func_nb_mt, a, b)
我突出显示了我更改的部分，其他所有内容都是从示例中逐字复制的。这利用了我机器上的所有内核(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 秒(对于 SIZE = 2147483648 * 1，在具有 16 个内核的机器上 32 个线程(。
更复杂的内核通常可以通过传入fastmath=True来加快速度。