Иногда вам нужно написать неидиоматический код 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
И вот некоторые моменты использования %timeitмагии IPython :
>>> %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)
Это составляет 65-кратное ускорение в меньшем случае и 26-кратное ускорение в большем случае (по сравнению с медленным зацикливанием кода, конечно) с использованием ускоренного кода. Другим преимуществом является то, что (в отличие от типичной векторизации с нативной нативой) нам не требовалась дополнительная память для достижения этой скорости, все дело в оптимизированном и скомпилированном низкоуровневом коде, который в конечном итоге запускается.
Вышеприведенная функция предполагает, что int64по умолчанию используются пустые массивы int , чего нет в Windows. Таким образом, альтернативой является удаление подписи из вызова to numba.njit, запускающей правильную сборку точно в срок. Но это означает, что функция будет скомпилирована во время первого выполнения, что может повлиять на результаты синхронизации (мы можем либо выполнить функцию один раз вручную, используя репрезентативные типы данных, либо просто принять, что первое выполнение синхронизации будет намного медленнее, что должно быть проигнорированным). Это именно то, что я пытался предотвратить, указав сигнатуру, которая запускает преждевременную компиляцию.
Во всяком случае, в надлежащем случае JIT декоратор нам нужен просто
@numba.njit
def diffmedian_jit(...):
Обратите внимание, что приведенные выше моменты времени, которые я показал для функции jit-compiled, применяются только после того, как функция была скомпилирована. Это может происходить либо при определении (с готовой компиляцией, когда передается явная подпись numba.njit), либо во время первого вызова функции (с отложенной компиляцией, когда подпись не передается numba.njit). Если функция будет выполняться только один раз, то время компиляции также должно учитываться для скорости этого метода. Как правило, компиляция функций имеет смысл только в том случае, если общее время выполнения + компиляции меньше, чем некомпилированное время выполнения (что на самом деле верно в вышеупомянутом случае, когда собственная функция python очень медленная). В основном это происходит, когда вы часто вызываете скомпилированную функцию.
Как max9111 отмечена в комментариях, одна важная особенностью numbaявляется cacheключевым словом , чтобы jit. Передача cache=Trueto numba.jitсохранит скомпилированную функцию на диск, так что во время следующего выполнения данного модуля python функция будет загружена оттуда, а не перекомпилирована, что снова может сэкономить вам время выполнения в долгосрочной перспективе.
scipy.ndimage.medianпредложение в связанном ответе? Мне не кажется, что для каждого ярлыка нужно одинаковое количество элементов. Или я что-то пропустил?