Почему этот код numba в 6 раз медленнее, чем код numpy?

Ответ 1

Странно, что numba может быть намного медленнее.

Это не слишком странно. Когда вы вызываете функции NumPy внутри функции numba, вы вызываете numba-версию этих функций. Это может быть быстрее, медленнее или так же быстро, как версии NumPy. Возможно, вам повезет, или вам не повезло (вам не повезло!). Но даже в функции numba вы все еще создаете много временных времен, потому что вы используете функции NumPy (один временной массив для результата точки, по одному для каждого квадрата и суммы, один для точки плюс первая сумма), поэтому вы не пользуетесь преимуществами возможности с numba.

Я использую это неправильно?

По существу: Да.

Мне действительно нужно ускорить это

Хорошо, я попробую.

Начните с разворачивания суммы квадратов по осям 1:

import numba as nb

@nb.njit
def sum_squares_2d_array_along_axis1(arr):
    res = np.empty(arr.shape[0], dtype=arr.dtype)
    for o_idx in range(arr.shape[0]):
        sum_ = 0
        for i_idx in range(arr.shape[1]):
            sum_ += arr[o_idx, i_idx] * arr[o_idx, i_idx]
        res[o_idx] = sum_
    return res


@nb.njit
def euclidean_distance_square_numba_v1(x1, x2):
    return -2 * np.dot(x1, x2.T) + np.expand_dims(sum_squares_2d_array_along_axis1(x1), axis=1) + sum_squares_2d_array_along_axis1(x2)

На моем компьютере это уже в 2 раза быстрее, чем код NumPy, и почти в 10 раз быстрее, чем исходный код Numba.

Говоря из опыта, получая его в 2 раза быстрее, чем NumPy, как правило, это ограничение (по крайней мере, если версия NumPy не является излишне сложной или неэффективной), однако вы можете выжать еще немного, развернув все:

import numba as nb

@nb.njit
def euclidean_distance_square_numba_v2(x1, x2):
    f1 = 0.
    for i_idx in range(x1.shape[1]):
        f1 += x1[0, i_idx] * x1[0, i_idx]

    res = np.empty(x2.shape[0], dtype=x2.dtype)
    for o_idx in range(x2.shape[0]):
        val = 0
        for i_idx in range(x2.shape[1]):
            val_from_x2 = x2[o_idx, i_idx]
            val += (-2) * x1[0, i_idx] * val_from_x2 + val_from_x2 * val_from_x2
        val += f1
        res[o_idx] = val
    return res

Но это только дает улучшение на 10-20% по сравнению с последним подходом.

В этот момент вы можете понять, что можете упростить код (хотя он, вероятно, не ускорит его):

import numba as nb

@nb.njit
def euclidean_distance_square_numba_v3(x1, x2):
    res = np.empty(x2.shape[0], dtype=x2.dtype)
    for o_idx in range(x2.shape[0]):
        val = 0
        for i_idx in range(x2.shape[1]):
            tmp = x1[0, i_idx] - x2[o_idx, i_idx]
            val += tmp * tmp
        res[o_idx] = val
    return res

Да, это выглядит довольно прямолинейно, и это не так уж медленно.

Однако во всем волнении я забыл упомянуть очевидное решение: scipy.spatial.distance.cdist которого есть sqeuclidean (квадрат евклидова расстояния):

from scipy.spatial import distance
distance.cdist(x1, x2, metric='sqeuclidean')

Это не намного быстрее, чем numba, но доступно без необходимости писать свою собственную функцию...

тесты

Проверьте правильность и сделайте разминки:

x1 = np.array([[1.,2,3]])
x2 = np.array([[1.,2,3], [2,3,4], [3,4,5], [4,5,6], [5,6,7]])

res1 = euclidean_distance_square(x1, x2)
res2 = euclidean_distance_square_numba_original(x1, x2)
res3 = euclidean_distance_square_numba_v1(x1, x2)
res4 = euclidean_distance_square_numba_v2(x1, x2)
res5 = euclidean_distance_square_numba_v3(x1, x2)
np.testing.assert_array_equal(res1, res2)
np.testing.assert_array_equal(res1, res3)
np.testing.assert_array_equal(res1[0], res4)
np.testing.assert_array_equal(res1[0], res5)
np.testing.assert_almost_equal(res1, distance.cdist(x1, x2, metric='sqeuclidean'))

Тайминги:

x1 = np.random.random((1, 512))
x2 = np.random.random((1000000, 512))

%timeit euclidean_distance_square(x1, x2)
# 2.09 s ± 54.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit euclidean_distance_square_numba_original(x1, x2)
# 10.9 s ± 158 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit euclidean_distance_square_numba_v1(x1, x2)
# 907 ms ± 7.11 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit euclidean_distance_square_numba_v2(x1, x2)
# 715 ms ± 15 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit euclidean_distance_square_numba_v3(x1, x2)
# 731 ms ± 34.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit distance.cdist(x1, x2, metric='sqeuclidean')
# 706 ms ± 4.99 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Примечание. Если у вас есть массивы целых чисел, вы можете захотеть изменить hardcoded 0.0 в numba-функциях на 0.

Ответ 2

Это комментарий к ответу @MSeifert. Есть еще кое-что, чтобы повысить производительность. Как и в каждом цифровом коде, рекомендуется подумать о том, какой тип данных достаточно точен для вашей проблемы. Часто float32 также достаточно, иногда даже float64 недостаточно.

Я хочу также упомянуть здесь ключевое слово fastmath, которое может дать здесь 1,7-кратную скорость.

[Редактировать]

Для простого суммирования я изучил LLVM-код и обнаружил, что суммирование было разделено частичными суммами на векторию. (4 частичных суммы для double и 8 для float с использованием AVX2). Это необходимо исследовать дальше.

Код

import llvmlite.binding as llvm
llvm.set_option('', '--debug-only=loop-vectorize')

@nb.njit
def euclidean_distance_square_numba_v3(x1, x2):
    res = np.empty(x2.shape[0], dtype=x2.dtype)
    for o_idx in range(x2.shape[0]):
        val = 0
        for i_idx in range(x2.shape[1]):
            tmp = x1[0, i_idx] - x2[o_idx, i_idx]
            val += tmp * tmp
        res[o_idx] = val
    return res

@nb.njit(fastmath=True)
def euclidean_distance_square_numba_v4(x1, x2):
    res = np.empty(x2.shape[0], dtype=x2.dtype)
    for o_idx in range(x2.shape[0]):
        val = 0.
        for i_idx in range(x2.shape[1]):
            tmp = x1[0, i_idx] - x2[o_idx, i_idx]
            val += tmp * tmp
        res[o_idx] = val
    return res

@nb.njit(fastmath=True,parallel=True)
def euclidean_distance_square_numba_v5(x1, x2):
    res = np.empty(x2.shape[0], dtype=x2.dtype)
    for o_idx in nb.prange(x2.shape[0]):
        val = 0.
        for i_idx in range(x2.shape[1]):
            tmp = x1[0, i_idx] - x2[o_idx, i_idx]
            val += tmp * tmp
        res[o_idx] = val
    return res

Задержки

float64
x1 = np.random.random((1, 512))
x2 = np.random.random((1000000, 512))

0.42 v3 @MSeifert
0.25 v4
0.18 v5 parallel-version
0.48 distance.cdist

float32
x1 = np.random.random((1, 512)).astype(np.float32)
x2 = np.random.random((1000000, 512)).astype(np.float32)

0.09 v5

Как явно объявлять типы

В общем, я бы не рекомендовал этого. Ваши входные массивы могут быть C-contigous (как testdata) Fortran contigous или strided. Если вы знаете, что ваши данные всегда являются C-contiguos, вы можете написать

@nb.njit('double[:](double[:, ::1],double[:, ::1])',fastmath=True)
def euclidean_distance_square_numba_v6(x1, x2):
    res = np.empty(x2.shape[0], dtype=x2.dtype)
    for o_idx in range(x2.shape[0]):
        val = 0.
        for i_idx in range(x2.shape[1]):
            tmp = x1[0, i_idx] - x2[o_idx, i_idx]
            val += tmp * tmp
        res[o_idx] = val
    return res

Это обеспечивает такую же производительность, что и версия v4, но будет терпеть неудачу, если входные массивы не являются C-contigous или нет dtype = np.float64.

Вы также можете использовать

@nb.njit('double[:](double[:, :],double[:, :])',fastmath=True)
def euclidean_distance_square_numba_v7(x1, x2):
    res = np.empty(x2.shape[0], dtype=x2.dtype)
    for o_idx in range(x2.shape[0]):
        val = 0.
        for i_idx in range(x2.shape[1]):
            tmp = x1[0, i_idx] - x2[o_idx, i_idx]
            val += tmp * tmp
        res[o_idx] = val
    return res

Это также будет работать на массивах с чередованием, но будет намного медленнее, чем версия выше на C-смежных массивах. (0,66 с против 0,25 с). Обратите также внимание на то, что ваша проблема довольно ограничена пропускной способностью памяти. Разница может быть выше при вычислениях с ЦП.

Если вы разрешите делать Numba задание для вас, оно будет автоматически обнаружено, если массив является смежным или нет (предоставление сложных входных данных при первой попытке, а не нескончаемые данные, приведет к перекомпиляции)

Ответ 3

Несмотря на то, что ответ @MSeifert делает этот ответ довольно устаревшим, я все еще его размещаю, потому что он объясняет более подробно, почему numba-версия была медленнее, чем numpy-версия.

Как мы увидим, основным виновником являются различные шаблоны доступа к памяти для numpy и numba.

Мы можем воспроизвести поведение с гораздо более простой функцией:

import numpy as np
import numba as nb

def just_sum(x2):
    return np.sum(x2, axis=1)

@nb.jit('double[:](double[:, :])', nopython=True)
def nb_just_sum(x2):
    return np.sum(x2, axis=1)

x2=np.random.random((2048,2048))

И теперь тайминги:

>>> %timeit just_sum(x)
2.33 ms ± 71.9 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
>>> %timeit nb_just_sum(x)
33.7 ms ± 296 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

это означает, что numpy примерно в 15 раз быстрее!

При компиляции кода numba с аннотациями (например, numba --annotate-html sum.html numba_sum.py) мы видим, как сумма выполняется numba (см. Полный список суммирования в приложении):

1. инициализировать столбец результатов
2. добавьте весь первый столбец в столбец результатов
3. добавьте весь второй столбец в столбец результатов
4. и так далее

В чем проблема такого подхода? Макет памяти! Массив хранится в строчном порядке, и, таким образом, его чтение по столбцам приводит к гораздо большему количеству пропусков кеша, чем чтение по ряду (это то, что делает numpy). Существует отличная статья, которая объясняет возможные эффекты кеша.

Как мы видим, сумма-реализация numba еще не очень зрелая. Однако из рассмотренного рассмотрения реализация numba может быть конкурентоспособной для крупномасштабного столбца (т.е. Транспонированной матрицы):

>>> %timeit just_sum(x.T)
3.09 ms ± 66.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
>>> %timeit nb_just_sum(x.T)
3.58 ms ± 45.8 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

и это действительно так.

Как показал код @MSeifert, основным преимуществом numba является то, что с его помощью мы можем уменьшить количество временных массивов numpy. Однако некоторые вещи, которые выглядят легко, нелегки и наивное решение может быть довольно плохим. Построение суммы - такая операция - не следует думать, что простой цикл достаточно хорош - см. например, этот вопрос.

Листинг numba-sumation:

 Function name: array_sum_impl_axis
in file: /home/ed/anaconda3/lib/python3.6/site-packages/numba/targets/arraymath.py
with signature: (array(float64, 2d, A), int64) -> array(float64, 1d, C)
show numba IR
194:    def array_sum_impl_axis(arr, axis):
195:        ndim = arr.ndim
196:    
197:        if not is_axis_const:
198:            # Catch where axis is negative or greater than 3.
199:            if axis < 0 or axis > 3:
200:                raise ValueError("Numba does not support sum with axis"
201:                                 "parameter outside the range 0 to 3.")
202:    
203:        # Catch the case where the user misspecifies the axis to be
204:        # more than the number of the array dimensions.
205:        if axis >= ndim:
206:            raise ValueError("axis is out of bounds for array")
207:    
208:        # Convert the shape of the input array to a list.
209:        ashape = list(arr.shape)
210:        # Get the length of the axis dimension.
211:        axis_len = ashape[axis]
212:        # Remove the axis dimension from the list of dimensional lengths.
213:        ashape.pop(axis)
214:        # Convert this shape list back to a tuple using above intrinsic.
215:        ashape_without_axis = _create_tuple_result_shape(ashape, arr.shape)
216:        # Tuple needed here to create output array with correct size.
217:        result = np.full(ashape_without_axis, zero, type(zero))
218:    
219:        # Iterate through the axis dimension.
220:        for axis_index in range(axis_len):
221:            if is_axis_const:
222:                # constant specialized version works for any valid axis value
223:                index_tuple_generic = _gen_index_tuple(arr.shape, axis_index,
224:                                                       const_axis_val)
225:                result += arr[index_tuple_generic]
226:            else:
227:                # Generate a tuple used to index the input array.
228:                # The tuple is ":" in all dimensions except the axis
229:                # dimension where it is "axis_index".
230:                if axis == 0:
231:                    index_tuple1 = _gen_index_tuple(arr.shape, axis_index, 0)
232:                    result += arr[index_tuple1]
233:                elif axis == 1:
234:                    index_tuple2 = _gen_index_tuple(arr.shape, axis_index, 1)
235:                    result += arr[index_tuple2]
236:                elif axis == 2:
237:                    index_tuple3 = _gen_index_tuple(arr.shape, axis_index, 2)
238:                    result += arr[index_tuple3]
239:                elif axis == 3:
240:                    index_tuple4 = _gen_index_tuple(arr.shape, axis_index, 3)
241:                    result += arr[index_tuple4]
242:    
243:        return result