Numba `nogil` + dask threading backend не приводит к ускорению (вычисления медленнее!) - PullRequest
Новая
       16

Numba `nogil` + dask threading backend не приводит к ускорению (вычисления медленнее!)

2 голосов
/ 02 июля 2019

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

К сожалению, Dask не приведет к ускорению, если я не использую планировщик 'processes'. Если вместо этого я использую 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

Как для обработки, так и для потокового бэкенда, я вижу полное использование всех ядер ЦП, как и ожидалось. Но не надо ускоряться для потокового бэкэнда. Я почти уверен, что функция jited, jit_render_internal, на самом деле не освобождает GIL.

Мои два вопроса:

  1. Если ключевое слово nogil передано numba.jit и GIL не может быть освобожден, почему не возникает ошибка?
  2. Почему код, как я его написал, не выпускает GIL? Все вычисления встроены в функцию и не имеют возвращаемого значения.

1 Ответ

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

Попробуйте следующее, что намного быстрее и, похоже, решает проблему производительности потока: 

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 на выше. На моей машине решение, основанное на потоках, на самом деле быстрее, чем процессы после изменения.

...