Numba `nogil` + dask 线程后端导致没有加速(计算更慢!)
Numba `nogil` + dask threading backend results in no speed up (computation is slower!)
我正在尝试使用 Numba and Dask to speed up a slow computation that is similar to calculating the kernel density estimate 大量积分。我的计划是在 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 核心的完全利用,正如预期的那样。但是线程后端没有加速。我很确定 jitted 函数 jit_render_internal
实际上并没有释放 GIL。
我的两个问题是:
- 如果
nogil
关键字传递给numba.jit
,GIL无法释放,为什么不报错?
- 为什么我编写的代码没有发布 GIL?所有计算都嵌入函数中,没有 return 值。
尝试以下方法,速度更快并且似乎解决了线程性能问题:
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
我在上面拆分了 x0
和 y0
的计算。在我的机器上,基于线程的解决方案实际上比更改后的进程更快。
我正在尝试使用 Numba and Dask to speed up a slow computation that is similar to calculating the kernel density estimate 大量积分。我的计划是在 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 核心的完全利用,正如预期的那样。但是线程后端没有加速。我很确定 jitted 函数 jit_render_internal
实际上并没有释放 GIL。
我的两个问题是:
- 如果
nogil
关键字传递给numba.jit
,GIL无法释放,为什么不报错? - 为什么我编写的代码没有发布 GIL?所有计算都嵌入函数中,没有 return 值。
尝试以下方法,速度更快并且似乎解决了线程性能问题:
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
我在上面拆分了 x0
和 y0
的计算。在我的机器上,基于线程的解决方案实际上比更改后的进程更快。