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

Есть ли причина, по которой следующий код запускается за 2 секунды,

def euclidean_distance_square(x1, x2):
    return -2*np.dot(x1, x2.T) + np.expand_dims(np.sum(np.square(x1), axis=1), axis=1) + np.sum(np.square(x2), axis=1)

а следующий код numba выполняется за 12 секунд?

@jit(nopython=True)
def euclidean_distance_square(x1, x2):
   return -2*np.dot(x1, x2.T) + np.expand_dims(np.sum(np.square(x1), axis=1), axis=1) + np.sum(np.square(x2), axis=1)

Мой x1 является матрицей размерности (1, 512), а x2 - матрицей размерности (3000000, 512). Странно, что numba может быть намного медленнее. Я использую это неправильно?

Мне действительно нужно ускорить это, потому что мне нужно запустить эту функцию 3 миллиона раз, а 2s все еще слишком медленный.

Мне нужно запустить это на CPU, потому что, поскольку вы можете видеть, что размер x2 настолько огромен, его нельзя загрузить на графический процессор (или, по крайней мере, мой GPU), недостаточно памяти.

Ответ 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