Почему этот код Numba в 6 раз медленнее, чем код NumPy?
Есть ли причина, по которой следующий код выполняется в 2s,
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). Довольно странно, что нумба может быть намного медленнее. Я использую это неправильно?
Мне действительно нужно ускорить это, потому что мне нужно запустить эту функцию 3 миллиона раз, а 2s все еще слишком медленно.
Мне нужно запустить это на процессоре, потому что, как вы видите, размерность x2 настолько велика, что ее нельзя загрузить на графический процессор (или, по крайней мере, на мой графический процессор), не хватает памяти.
3 ответа
Довольно странно, что нумба может быть намного медленнее.
Это не слишком странно. Когда вы вызываете функции 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)
Примечание: если у вас есть массивы целых чисел, вы можете изменить жестко запрограммированный 0.0 в функции Нумба в 0,
Несмотря на то, что ответ @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 (см. полный список суммирования в приложении):
- инициализировать столбец результата
- добавить весь первый столбец в столбец результата
- добавить весь второй столбец в столбец результата
- и так далее
В чем проблема этого подхода? Расположение памяти! Массив хранится в главном порядке строк, и, следовательно, чтение его по столбцам приводит к гораздо большему количеству промахов в кэше, чем чтение по строкам (что и делает 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-суммирование:
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's 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
Это комментарий к ответу @MSeifert. Есть еще несколько вещей, чтобы получить производительность. Как и в любом числовом коде, рекомендуется подумать о том, какой тип данных является достаточно точным для вашей задачи. Часто float32 также достаточно, иногда даже float64 недостаточно.
Я также хочу упомянуть ключевое слово fastmath, которое может увеличить скорость в 1,7 раза.
[Редактировать]
Для простого суммирования я заглянул в LLVM-код и обнаружил, что суммация была разделена на частичные суммы при векторизации. (4 частичные суммы для двойной и 8 для плавающей с использованием 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 (в качестве тестовых данных) Fortran, непрерывными или непрерывными. Если вы знаете, что ваши данные всегда 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-contigous массивов. (0,66 с против 0,25 с). Обратите внимание также, что ваша проблема довольно ограничена пропускной способностью памяти. Разница может быть выше при вычислениях с привязкой к процессору.
Если вы дадите Numba работу за вас, она будет автоматически обнаружена, если массив является непрерывным или нет (предоставление непрерывных входных данных с первой попытки, а затем не непрерывных данных приведет к перекомпиляции)