numba`nogil` dask螺纹后端导致无速度(计算较慢!)



我正在尝试使用numba和dask加快慢速计算,该计算类似于计算大量点的内核密度估计。我的计划是在jit ED函数中编写计算上昂贵的逻辑,然后使用dask在CPU内核中分配工作。我想使用numba.jit函数的nogil功能,以便可以使用dask线程后端,以避免输入数据的不必要的内存副本(这很大(。

不幸的是,除非我使用'processes'调度程序,否则DASK不会导致加速。如果我使用ThreadPoolExector,则看到预期的速度。

这是我问题的简化示例:

import os
import numpy as np
import numba
import dask
CPU_COUNT = os.cpu_count()
def render_internal(size, mag):
    """mag is the magnification to apply
    generate coordinates internally
    """
    coords = np.random.rand(size, 2)
    img = np.zeros((mag, mag), dtype=np.int64)
    for i in range(len(coords)):
        y0, x0 = coords[i] * mag
        y1, x1 = int(y0), int(x0)
        m = 1
        img[y1, x1] += m
jit_render_internal = numba.jit(render_internal, nogil=True, nopython=True)
args = 10000000, 100
print("Linear time:")
%time linear_compute = [jit_render_internal(*args) for i in range(CPU_COUNT)]
delayed_jit_render_internal = dask.delayed(jit_render_internal)
print()
print("Threads time:")
%time dask_compute_threads = dask.compute(*[delayed_jit_render_internal(*args) for i in range(CPU_COUNT)])
print()
print("Processes time:")
%time dask_compute_processes = dask.compute(*[delayed_jit_render_internal(*args) for i in range(CPU_COUNT)], scheduler="processes")

这是我的机器上的输出:

Linear time:
Wall time: 1min 17s
Threads time:
Wall time: 1min 47s
Processes time:
Wall time: 7.79 s

对于处理和螺纹后端,我看到了所有CPU内核的完整利用。但是没有加快螺纹后端的速度。我很确定奇特的函数jit_render_internal实际上不是释放GIL。

我的两个问题是:

  1. 如果nogil关键字传递给numba.jit并且无法释放GIL,为什么不引起错误?
  2. 为什么我写的代码不释放GIL?所有计算都嵌入到函数中,没有返回值。

尝试以下内容,这要快得多,并且似乎可以解决线程性能问题:

def render_internal(size, mag):
    """mag is the magnification to apply
    generate coordinates internally
    """
    coords = np.random.rand(size, 2)
    img = np.zeros((mag, mag), dtype=np.int64)
    for i in range(len(coords)):
        #y0, x0 = coords[i] * mag
        y0 = coords[i,0] * mag
        x0 = coords[i,1] * mag
        y1, x1 = int(y0), int(x0)
        m = 1
        img[y1, x1] += m

我已经在上面的x0y0的计算中分配了计算。在我的计算机上,基于线程的解决方案实际上比更改后的过程要快。

最新更新