Иногда вам нужно написать неидиоматический код numpy, если вы действительно хотите ускорить ваши вычисления, чего вы не можете сделать с собственным numpy.
numba компилирует ваш код Python в низкоуровневый C. Поскольку большая часть numpy обычно работает так же быстро, как C, это в основном оказывается полезным, если ваша проблема не поддается нативной векторизации с numpy. Это один из примеров (где я предположил, что индексы являются смежными и отсортированными, что также отражено в данных примера):
import numpy as np
import numba
# use the inflated example of roganjosh https://stackoverflow.com/a/58788534
data = [1.00, 1.05, 1.30, 1.20, 1.06, 1.54, 1.33, 1.87, 1.67]
index = [0, 0, 1, 1, 1, 1, 2, 3, 3]
data = np.array(data * 500) # using arrays is important for numba!
index = np.sort(np.random.randint(0, 30, 4500))
# jit-decorate; original is available as .py_func attribute
@numba.njit('f8[:](f8[:], i8[:])') # explicit signature implies ahead-of-time compile
def diffmedian_jit(data, index):
res = np.empty_like(data)
i_start = 0
for i in range(1, index.size):
if index[i] == index[i_start]:
continue
# here: i is the first _next_ index
inds = slice(i_start, i) # i_start:i slice
res[inds] = data[inds] - np.median(data[inds])
i_start = i
# also fix last label
res[i_start:] = data[i_start:] - np.median(data[i_start:])
return res
А вот некоторые временные примеры использования IPython %timeit magic:
>>> %timeit diffmedian_jit.py_func(data, index) # non-jitted function
... %timeit diffmedian_jit(data, index) # jitted function
...
4.27 ms ± 109 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
65.2 µs ± 1.01 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
Используя обновленные данные примера в вопросе, эти числа (т. Е. Время выполнения функции python по сравнению с временем выполнения JIT-ускоренной функции) равны
>>> %timeit diffmedian_jit.py_func(data, groups)
... %timeit diffmedian_jit(data, groups)
2.45 s ± 34.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
93.6 ms ± 518 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
Это составляет 65xускорение в меньшем случае и 26-кратное ускорение в большем случае (по сравнению с медленным цикличным кодом, конечно) с использованием ускоренного кода. Другим положительным моментом является то, что (в отличие от типичной векторизации с нативным Numpy) нам не требовалась дополнительная память для достижения этой скорости, речь идет об оптимизированном и скомпилированном низкоуровневом коде, который в конечном итоге запускается.
Приведенная выше функция предполагает, что numpy int-массивы по умолчанию int64, что на самом деле не так в Windows. Таким образом, альтернативой является удаление подписи из звонка на numba.njit, что вызывает правильную компиляцию точно в срок. Но это означает, что функция будет скомпилирована во время первого выполнения, что может повлиять на результаты синхронизации (мы можем либо выполнить функцию один раз вручную, используя репрезентативные типы данных, либо просто принять, что первое выполнение синхронизации будет намного медленнее, что должнобыть проигнорированным). Это именно то, что я пытался предотвратить, указав сигнатуру, которая запускает преждевременную компиляцию.
В любом случае, в правильном JIT-случае декоратор, который нам нужен, это просто
@numba.njit
def diffmedian_jit(...):
Обратите внимание, что приведенные выше моменты времени, которые я показал для функции jit-compiled, применяются только после того, как функция была скомпилирована. Это может происходить либо при определении (с готовой компиляцией, когда явная подпись передается numba.njit), либо во время первого вызова функции (с отложенной компиляцией, когда никакая подпись не передается numba.njit). Если функция будет выполняться только один раз, время компиляции также должно учитываться для скорости этого метода. Как правило, компиляция функций имеет смысл только в том случае, если общее время выполнения + компиляции меньше, чем некомпилированное время выполнения (что на самом деле верно в вышеупомянутом случае, когда собственная функция python очень медленная). В основном это происходит, когда вы часто вызываете скомпилированную функцию.
Как отмечается в комментарии max9111 , одной важной особенностью numba является ключевое слово cache. до jit. Передача cache=True в numba.jit сохранит скомпилированную функцию на диск, так что во время следующего выполнения данного модуля python функция будет загружена оттуда, а не перекомпилирована, что снова может сэкономить вам время выполнения в долгосрочной перспективе.