Блог

Я пытаюсь ускорить фрагмент кода, свертывающий 1-мерный массив (фильтр) по каждому столбцу 2D-массива. Каким-то образом, когда я запускаю его с помощью numba njit , я получаю замедление в 7 раз. Мои мысли:

• Возможно, индексация столбцов замедляет ее, но переключение на индексацию строк не повлияло на производительность
• Возможно, индексация результатов свертки происходит медленно, но ее удаление ничего не изменило
• Я проверил, что numba правильно понимает все типы

(протестировано на Windows 10, python 3.9.4 от conda, numpy 1.12.2, numba 0.53.1)

Кто-нибудь может сказать мне, почему этот код медленный?

 import numpy as np
from numba import njit

def f1(a1, filt):
    l2 = filt.size // 2
    res = np.empty(a1.shape)
    for i in range(a1.shape[1]):
        res[:, i] = np.convolve(a1[:, i], filt)[l2:-l2]
    return res

@njit
def f1_jit(a1, filt):
    l2 = filt.size // 2
    res = np.empty(a1.shape)
    for i in range(a1.shape[1]):
        res[:, i] = np.convolve(a1[:, i], filt)[l2:-l2]
    return res

a1 = np.random.random((6400, 1000))
filt = np.random.random((65))
f1(a1, filt)
f1_jit(a1, filt)

%timeit f1(a1, filt)     # 404 ms ± 19.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit f1_jit(a1, filt) # 2.8 s ± 66.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
 

Проблема возникает из-за реализации Numba np.convolve . Это известная проблема. Оказывается, текущая реализация Numba намного медленнее, чем реализация Numpy (версия <= 0.54.1, протестированная в Windows).

Под капотом

С одной стороны, вызов реализации Numpy correlate , который сам выполняет скалярное произведение, которое должно быть реализовано библиотекой fast BLAS, доступной в вашей системе. С другой стороны, вызовы реализации Numba _get_inner_prod , которые используют np.dot это, также должны использовать ту же библиотеку BLAS (при условии, что обнаружен BLAS, что должно быть так)…

При этом существует множество проблем, связанных с точечным произведением:

Прежде всего, если внутренняя переменная _HAVE_BLAS numba/np/arraymath.py отключена вручную, Numba использует резервную реализацию точечного произведения, которая должна быть значительно медленнее. Однако оказывается, что использование резервной реализации точечного продукта, используемой np.convolve результатом, в 5 раз быстрее, чем с оболочкой BLAS на моей машине! Использование дополнительного параметра fastmath=True в njit декораторе Numba приводит к общему ускорению выполнения в 8,7 раза! Вот тестовый код:

 import numpy as np
import numba as nb

def npConvolve(a, b):
    return np.convolve(a, b)

@nb.njit('float64[:](float64[:], float64[:])')
def nbConvolveUncont(a, b):
    return np.convolve(a, b)

@nb.njit('float64[::1](float64[::1], float64[::1])')
def nbConvolveCont(a, b):
    return np.convolve(a, b)

a = np.random.random(6400)
b = np.random.random(65)
%timeit -n 100 npConvolve(a, b)
%timeit -n 100 nbConvolveUncont(a, b)
%timeit -n 100 nbConvolveCont(a, b)
 

Вот необработанные интересные результаты:

 With _HAVE_BLAS=True (default):
126 µs ± 292 ns per loop
1.6 ms ± 21.3 µs per loop
1.6 ms ± 18.5 µs per loop

With _HAVE_BLAS=False:
125 µs ± 359 ns per loop
311 µs ± 1.18 µs per loop
268 µs ± 4.26 µs per loop

With _HAVE_BLAS=False and fastmath=True:
125 µs ± 757 ns per loop
327 µs ± 3.69 µs per loop
183 µs ± 654 ns per loop
 

Более того, np_convolve Numba внутренне переворачивает некоторый параметр массива, а затем выполняет скалярное произведение, используя перевернутый массив, который имеет нетривиальный шаг (т. Е. не 1). Такой нетривиальный шаг может повлиять на производительность точечного продукта. В более общем плане, любое преобразование, не позволяющее компилятору узнать, что массивы являются смежными, безусловно, сильно повлияет на производительность. Действительно, следующий тест показывает влияние работы над непрерывным массивом с реализацией точечного произведения Numba:

 import numpy as np
import numba as nb

def np_dot(a, b):
    return np.dot(a, b)

@nb.njit('float64(float64[::1], float64[::1])')
def nb_dot_cont(a, b):
    return np.dot(a, b)

@nb.njit('float64(float64[::1], float64[:])')
def nb_dot_stride(a, b):
    return np.dot(a, b)

v = np.random.random(128*1024)
%timeit -n 200 np_dot(v, v)         #  36.5 µs ±  4.9 µs per loop
%timeit -n 200 nb_dot_stride(v, v)  # 361.0 µs ± 17.1 µs per loop  (x10 !!!)
%timeit -n 200 nb_dot_cont(v, v)    #  34.1 µs ±  2.9 µs per loop
 

Некоторые общие замечания о Numpy и Numba

Обратите внимание, что Numba вряд ли может ускорить вызовы Numpy, когда они работают с довольно большими массивами, поскольку Numba повторно реализует функции Numpy в основном на Python и использует JIT-компилятор (LLVM-Lite) для их ускорения, в то время как Numpy в основном реализован на обычном C (с довольно медленным кодом переноса Python). Код Numpy использует низкоуровневые функции процессора, такие как инструкции SIMD, чтобы значительно ускорить выполнение многих функций. Оба, похоже, используют библиотеки BLAS, которые, как известно, высоко оптимизированы. Numpy, как правило, более оптимизирован, поскольку Numpy в настоящее время более зрелый, чем Numba: у Numpy гораздо больше участников, работающих с более длительного времени.