Numpy `np.prod`, `np.all` работают медленнее, чем явные алгебраические операции

Ваш тестовый пример:

In [2]: N_ROWS = int(2e5)
   ...: N_ITER = int(1e3)
   ...: 
   ...: np.random.seed(123)
   ...: 
   ...: mat = np.random.rand(N_ROWS, 2)
   ...: mat_bool = mat > .5

prod против оператора, проверяя, что они производят одно и то же:

In [3]: x=np.prod(mat, axis=1); x.shape
Out[3]: (200000,)

In [4]: y=mat[:,0]*mat[:,1]; y.shape
Out[4]: (200000,)

In [5]: np.allclose(x,y)
Out[5]: True

Times - prod примерно в 5 раз медленнее:

In [6]: timeit x=np.prod(mat, axis=1)
5.84 ms ± 13.7 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [7]: timeit y=mat[:,0]*mat[:,1]
1.18 ms ± 38.9 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

Это версия numpy 1.24, относительно новая, хотя и не самая последняя.

Альтернативные тайминги показывают, что оператор использует multiplyufunc, а prod использует ufunc'sreduce. Это также видно из prod [исходного] кода.

In [8]: timeit np.multiply(mat[:,0],mat[:,1])
1.16 ms ± 18.5 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

In [9]: timeit np.multiply.reduce(mat, axis=1)
5.82 ms ± 20.3 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

np.all использует np.logical_andufunc таким же образом. Очевидно, за использованием ufunc.reduce стоит немного больше «багажа».

Основная проблема заключается в том, что реализация Numpy не рассчитана на эффективность, когда операции выполняются по очень маленькой оси.

Операции Numpy, повторяющиеся по очень маленькой оси, довольно неэффективны. Это связано с внутренними накладными расходами Numpy. Особенно внутренний итератор Numpy. Ufuncs также являются значительным источником накладных расходов в текущей реализации. Все это описано в этом посте . Часть накладных расходов можно сократить, но я не ожидаю значительного улучшения в ближайшее время. Другой невозможно уменьшить, не создавая серьезных проблем с удобством сопровождения в реализации Numpy.

Лучшее, что можно сделать (как указано в сообщении выше), — это работать с транспонированным макетом массива. Это не только значительно снижает накладные расходы Numpy (которые все еще немаловажны), но и делает операцию SIMD-дружественной, что в настоящее время имеет решающее значение для производительности.

Более того, вы можете использовать объектные методы вместо глобальных функций Numpy, чтобы немного уменьшить накладные расходы при запуске Numpy. Ставьте лайк mat_bool.all(axis=0) вместо np.all(mat_bool, axis=0). Это имеет смысл только в том случае, если массивы небольшие (в противном случае накладные расходы при запуске должны быть относительно небольшими).

Вот модифицированный код:

import numpy as np
import time

N_ROWS = int(2e5)
N_ITER = int(1e3)

np.random.seed(123)

mat = np.random.rand(N_ROWS, 2)

# Actual (physical) array transposition
# A copy is needed since mat.T is just a view with the same memory layout
mat = mat.T.copy()

mat_bool = mat > .5

# np.prod
start = time.time()
for i in range(N_ITER):
    _ = mat.prod(axis=0)
duration = time.time() - start
print(f"np.prod took {duration}s")

start = time.time()
for i in range(N_ITER):
    _ = mat[0, :] * mat[1, :]
duration = time.time() - start
print(f"Manual prod took {duration}s")

# np.all
start = time.time()
for i in range(N_ITER):
    _ = mat_bool.all(axis=0)
duration = time.time() - start
print(f"np.all took {duration}s")

start = time.time()
for i in range(N_ITER):
    _ = mat_bool[0, :] * mat_bool[1, :]
duration = time.time() - start
print(f"manual all took {duration}s")

Вот первоначальные результаты на моем процессоре i5-9600KF с Numpy 1.24.3 в Windows:

np.prod took 2.087920904159546s
Manual prod took 0.4368312358856201s
np.all took 3.4697189331054688s
manual all took 1.1419456005096436s

Вот результаты предоставленного кода:

np.prod took 0.5076425075531006s
Manual prod took 0.4059138298034668s
np.all took 0.03490638732910156s
manual all took 0.014960050582885742s

Вот ускорение между исходным кодом и предоставленным:

np.prod:       4.11
Manual prod:   1.08
np.all:       99.40
manual all:   76.33

Ускорение колоссальное. Выгода от выполнения ручных операций по-прежнему есть, но разрыв существенно меньше (соответственно 25% и 133% для предоставленного кода, в отличие от соответственно 378% и 204% для исходного кода).

Даже после транспонирования часть накладных расходов по-прежнему связана с ufuncs, как указывает hpaulj, но теперь накладные расходы почти приемлемы. Вот тайминги:

In [94]: %timeit np.multiply(mat[0,:],mat[1::])
400 µs ± 11.4 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

In [95]: %timeit np.multiply.reduce(mat, axis=0)
510 µs ± 5.23 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Что касается скорости логических операций, то они сейчас настолько быстрые, что время, затрачиваемое на накладные расходы Numpy, становится значительным. В этом масштабе (микросекунды) я не думаю, что стоит тратить время на понимание того, почему одна функция быстрее/медленнее другой, поскольку обе функции в любом случае будут медленными, а более быструю реализацию часто можно разработать благодаря инструментам векторизации и компилятора, таким как Numba/ Cython в данном конкретном случае.

Также обратите внимание, что np.logical_and лучше подходит, чем произведение для логических значений (и на моей машине он работает одинаково быстро, хотя теоретически может быть быстрее).