我使用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)
为了完整起见,在2018年(numba v 0.39),您可以
在原始函数定义中将
range
替换为prange
,就这样。这立即使CPU利用率达到100%,在我的例子中,运行时间从2.9秒提高到1.7秒(对于大小为2147483648*1的机器,16核32线程)。
更复杂的内核通常可以通过传入
fastmath=True
来加快速度。您可以将
parallel=True
传递给任何numba jitted函数,但这并不意味着它总是使用所有核心。您必须理解numba使用一些启发式方法使代码并行执行,有时这些启发式方法在代码中根本找不到任何可并行化的东西。当前有一个pull request,因此如果无法使其“并行”,它会发出警告。所以它更像是一个“请尽可能让它并行执行”参数,而不是“强制并行执行”。但是,如果您真的知道可以并行化代码,则始终可以手动使用线程或进程。只是调整一下example of using multi-threading from the numba docs:
我强调了我修改过的部分,其他部分都是从示例中一字不差地复制出来的。这利用了我的机器上的所有核心(4核心机器,因此4个线程),但没有显示出显著的加速:
在这种情况下,多线程没有(很多)加速,这是因为添加是一个带宽受限的操作。这意味着从数组中加载元素并将结果放入结果数组所需的时间比实际添加所需的时间要长得多。
在这些情况下,您甚至可以看到由于并行执行而导致的减速!
只有当函数更复杂并且实际操作比加载和存储数组元素花费大量时间时,并行执行才会有很大的改进。numba文档中的示例如下:
这实际上(几乎)随线程数而变化,因为两个乘法、一个加法和一个调用
math.exp
比加载和存储结果慢得多:结果:
相关问题 更多 >
编程相关推荐