Numba `nogil = True` + ThreadPoolExecutor приводит к меньшей скорости, чем ожидалось - PullRequest
Новая
       26

Numba `nogil = True` + ThreadPoolExecutor приводит к меньшей скорости, чем ожидалось

0 голосов
/ 08 июля 2019

Это продолжение моего предыдущего вопроса :

Я пытаюсь использовать Numba и Dask для ускорения медленных вычислений, аналогичных вычислению оценки плотности ядра из огромной коллекции точки. Мой план состоял в том, чтобы написать вычислительно дорогую логику в jit ed-функции, а затем распределить работу между ядрами ЦП, используя dask. Я хотел использовать функцию nogil функции numba.jit, чтобы можно было использовать потоковый бэкэнд dask, чтобы избежать ненужных копий входных данных (которые очень велики).

Оказывается , что часть проблемы связана с распаковкой аргументов в jit ed-функции. Тем не менее, даже с этим исправлением я почти не вижу ускорения следующего кода: 

import os
import time
import numpy as np
from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor, as_completed
from numba import njit, jit

CPU_COUNT = os.cpu_count()
print("CPU_COUNT", CPU_COUNT)

def PE(pool, func, args):
    with pool(max_workers=CPU_COUNT) as exc:
        fut = {exc.submit(func, *arg): i for i, arg in enumerate(args)}
        for f in as_completed(fut):
            f.result()

def render(params, mag):
    """Render gaussian peaks in small windows"""

    radius = 3

    for i in range(len(params)):
        y0 = params[i, 0] * mag
        x0 = params[i, 1] * mag
        sy = params[i, 2] * mag
        sx = params[i, 3] * mag

        # calculate the render window size
        wy = int(sy * radius * 2.0)
        wx = int(sx * radius * 2.0)

        # calculate the area in the image
        ystart = int(np.rint(y0)) - wy // 2
        yend = ystart + wy
        xstart = int(np.rint(x0)) - wx // 2
        xend = xstart + wx

        # adjust coordinates to window coordinates
        y1 = y0 - ystart
        x1 = x0 - xstart

        y = np.arange(wy)
        x = np.arange(wx)
        amp = 1 / (2 * np.pi * sy * sx)
        gy = np.exp(-((y - y0) / sy) ** 2 / 2)
        gx = np.exp(-((x - x0) / sx) ** 2 / 2)
        g = amp * np.outer(gy, gx)

jit_render = jit(render, nopython=True, nogil=True)

args = [(np.random.rand(1000000, 4) * (1, 1, 0.02, 0.02), 100) for i in range(CPU_COUNT)]

print("Single time:")
# %timeit render(*args[0])
%timeit jit_render(*args[0])

print()
print("Linear time:")
%time [jit_render(*a) for a in args]

print()
print("Threads time:")
%time PE(ThreadPoolExecutor, jit_render, args)

На моем 8-ядерном MacBook скорость увеличивается примерно в 2 раза 

Single time:
1.6 s ± 153 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Linear time:
CPU times: user 11.8 s, sys: 43.1 ms, total: 11.9 s
Wall time: 11.9 s

Threads time:
CPU times: user 45.4 s, sys: 125 ms, total: 45.5 s
Wall time: 6.29 s

На моем 24-ядерном Windows-боксе скорость увеличивается примерно в 1 раз: 

Single time:
1.91 s ± 105 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Linear time:
Wall time: 1min 30s

Threads time:
Wall time: 1min 4s
...