Почему numpy einsum быстрее, чем встроенные функции numpy?

Давайте начнем с трех массивов dtype=np.double. Сроки выполняются на процессоре Intel, используя numpy 1.7.1, скомпилированный с icc и связанный с intel mkl. Для проверки таймингов также использовался процессор AMD c numpy 1.6.1, скомпилированный с gcc без mkl. Обратите внимание, что шкала таймингов почти линейно зависит от размера системы и не связана с небольшими накладными расходами, вызванными операторами nump функций if. Эти различия будут отображаться в микросекундах не в миллисекундах:

arr_1D=np.arange(500,dtype=np.double)
large_arr_1D=np.arange(100000,dtype=np.double)
arr_2D=np.arange(500**2,dtype=np.double).reshape(500,500)
arr_3D=np.arange(500**3,dtype=np.double).reshape(500,500,500)

Сначала рассмотрим функцию np.sum:

np.all(np.sum(arr_3D)==np.einsum('ijk->',arr_3D))
True

%timeit np.sum(arr_3D)
10 loops, best of 3: 142 ms per loop

%timeit np.einsum('ijk->', arr_3D)
10 loops, best of 3: 70.2 ms per loop

Пауэрс:

np.allclose(arr_3D*arr_3D*arr_3D,np.einsum('ijk,ijk,ijk->ijk',arr_3D,arr_3D,arr_3D))
True

%timeit arr_3D*arr_3D*arr_3D
1 loops, best of 3: 1.32 s per loop

%timeit np.einsum('ijk,ijk,ijk->ijk', arr_3D, arr_3D, arr_3D)
1 loops, best of 3: 694 ms per loop

Внешнее изделие:

np.all(np.outer(arr_1D,arr_1D)==np.einsum('i,k->ik',arr_1D,arr_1D))
True

%timeit np.outer(arr_1D, arr_1D)
1000 loops, best of 3: 411 us per loop

%timeit np.einsum('i,k->ik', arr_1D, arr_1D)
1000 loops, best of 3: 245 us per loop

Все вышеперечисленное в два раза быстрее с np.einsum. Это должны быть яблоки для сравнения яблок, так как все специально относится к dtype=np.double. Я ожидал бы ускорения в такой операции:

np.allclose(np.sum(arr_2D*arr_3D),np.einsum('ij,oij->',arr_2D,arr_3D))
True

%timeit np.sum(arr_2D*arr_3D)
1 loops, best of 3: 813 ms per loop

%timeit np.einsum('ij,oij->', arr_2D, arr_3D)
10 loops, best of 3: 85.1 ms per loop

Einsum, по-видимому, по крайней мере в два раза быстрее для np.inner, np.outer, np.kron и np.sum независимо от выбора axes. Основное исключение - np.dot, поскольку оно вызывает DGEMM из библиотеки BLAS. Итак, почему np.einsum быстрее, чем другие эквивалентные функции numpy?

Случай DGEMM для полноты:

np.allclose(np.dot(arr_2D,arr_2D),np.einsum('ij,jk',arr_2D,arr_2D))
True

%timeit np.einsum('ij,jk',arr_2D,arr_2D)
10 loops, best of 3: 56.1 ms per loop

%timeit np.dot(arr_2D,arr_2D)
100 loops, best of 3: 5.17 ms per loop

Ведущая теория - комментарий @sebergs, что np.einsum может использовать SSE2, но numpy ufuncs не будет до numpy 1.8 (см. изменить журнал). Я считаю, что это правильный ответ, но он не смог его подтвердить. Некоторое ограниченное доказательство можно найти, изменив тип входного массива и наблюдая разницу в скорости и тот факт, что не все наблюдают одни и те же тенденции в таймингах.

Ответ 1

Теперь, когда освобождается numpy 1.8, где согласно документам все ufuncs должны использовать SSE2, я хотел бы дважды проверить, что комментарий Seberg о SSE2 действителен.

Для выполнения теста была создана новая установка python 2.7 - с помощью icc были скомпилированы numpy 1.7 и 1.8 с использованием стандартных опций ядра AMD opteron, работающего под управлением Ubuntu.

Это тестовый запуск как до, так и после обновления 1.8:

import numpy as np
import timeit

arr_1D=np.arange(5000,dtype=np.double)
arr_2D=np.arange(500**2,dtype=np.double).reshape(500,500)
arr_3D=np.arange(500**3,dtype=np.double).reshape(500,500,500)

print 'Summation test:'
print timeit.timeit('np.sum(arr_3D)',
                      'import numpy as np; from __main__ import arr_1D, arr_2D, arr_3D',
                      number=5)/5
print timeit.timeit('np.einsum("ijk->", arr_3D)',
                      'import numpy as np; from __main__ import arr_1D, arr_2D, arr_3D',
                      number=5)/5
print '----------------------\n'


print 'Power test:'
print timeit.timeit('arr_3D*arr_3D*arr_3D',
                      'import numpy as np; from __main__ import arr_1D, arr_2D, arr_3D',
                      number=5)/5
print timeit.timeit('np.einsum("ijk,ijk,ijk->ijk", arr_3D, arr_3D, arr_3D)',
                      'import numpy as np; from __main__ import arr_1D, arr_2D, arr_3D',
                      number=5)/5
print '----------------------\n'


print 'Outer test:'
print timeit.timeit('np.outer(arr_1D, arr_1D)',
                      'import numpy as np; from __main__ import arr_1D, arr_2D, arr_3D',
                      number=5)/5
print timeit.timeit('np.einsum("i,k->ik", arr_1D, arr_1D)',
                      'import numpy as np; from __main__ import arr_1D, arr_2D, arr_3D',
                      number=5)/5
print '----------------------\n'


print 'Einsum test:'
print timeit.timeit('np.sum(arr_2D*arr_3D)',
                      'import numpy as np; from __main__ import arr_1D, arr_2D, arr_3D',
                      number=5)/5
print timeit.timeit('np.einsum("ij,oij->", arr_2D, arr_3D)',
                      'import numpy as np; from __main__ import arr_1D, arr_2D, arr_3D',
                      number=5)/5
print '----------------------\n'

Numpy 1.7.1:

Summation test:
0.172988510132
0.0934836149216
----------------------

Power test:
1.93524689674
0.839519000053
----------------------

Outer test:
0.130380821228
0.121401786804
----------------------

Einsum test:
0.979052495956
0.126066613197

Numpy 1.8:

Summation test:
0.116551589966
0.0920487880707
----------------------

Power test:
1.23683619499
0.815982818604
----------------------

Outer test:
0.131808176041
0.127472200394
----------------------

Einsum test:
0.781750011444
0.129271841049

Я думаю, что это довольно убедительно, что SSE играет большую роль в разнице во времени, следует отметить, что повторение этих тестов очень быстро, всего лишь на 0,003 секунды. Оставшаяся разница должна быть рассмотрена в других ответах на этот вопрос.

Ответ 2

Во-первых, в списке numpy было много последних обсуждений. Например, см. http://numpy-discussion.10968.n7.nabble.com/poor-performance-of-sum-with-sub-machine-word-integer-types-td41.html http://numpy-discussion.10968.n7.nabble.com/odd-performance-of-sum-td3332.html

Некоторые из них сводятся к тому, что einsum является новым и, по-видимому, пытается улучшить выравнивание кеша и другие проблемы с доступом к памяти, в то время как многие из более старых функций numpy сосредоточены на легко переносимой реализации на сильно оптимизированной один. Я просто размышляю, хотя.

Однако некоторые из того, что вы делаете, - это не совсем сравнение яблок с яблоками.

В дополнение к тому, что уже сказал @Jamie, sum использует более подходящий аккумулятор для массивов

Например, sum более осторожен в проверке типа ввода и использовании соответствующего аккумулятора. Например, рассмотрим следующее:

In [1]: x = 255 * np.ones(100, dtype=np.uint8)

In [2]: x
Out[2]:
array([255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
       255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
       255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
       255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
       255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
       255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
       255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
       255, 255, 255, 255, 255, 255, 255, 255, 255], dtype=uint8)

Обратите внимание, что sum верен:

In [3]: x.sum()
Out[3]: 25500

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

In [4]: np.einsum('i->', x)
Out[4]: 156

Но если мы используем менее ограниченный dtype, мы все равно получим результат, который вы ожидаете:

In [5]: y = 255 * np.ones(100)

In [6]: np.einsum('i->', y)
Out[6]: 25500.0

Ответ 3

Я думаю, что эти тайминги объясняют, что происходит:

a = np.arange(1000, dtype=np.double)
%timeit np.einsum('i->', a)
100000 loops, best of 3: 3.32 us per loop
%timeit np.sum(a)
100000 loops, best of 3: 6.84 us per loop

a = np.arange(10000, dtype=np.double)
%timeit np.einsum('i->', a)
100000 loops, best of 3: 12.6 us per loop
%timeit np.sum(a)
100000 loops, best of 3: 16.5 us per loop

a = np.arange(100000, dtype=np.double)
%timeit np.einsum('i->', a)
10000 loops, best of 3: 103 us per loop
%timeit np.sum(a)
10000 loops, best of 3: 109 us per loop

Таким образом, при вызове np.sum over np.einsum у вас в основном есть почти постоянный 3us накладные расходы, поэтому они в основном работают так же быстро, но требуется немного больше времени для перехода. Почему это может быть? Мои деньги заключаются в следующем:

a = np.arange(1000, dtype=object)
%timeit np.einsum('i->', a)
Traceback (most recent call last):
...
TypeError: invalid data type for einsum
%timeit np.sum(a)
10000 loops, best of 3: 20.3 us per loop

Не уверен, что происходит в точности, но кажется, что np.einsum пропускает некоторые проверки, чтобы извлечь специфичные для типа функции для выполнения умножений и дополнений и напрямую идет с * и + для стандартных типов C только.

Многомерные случаи не различны:

n = 10; a = np.arange(n**3, dtype=np.double).reshape(n, n, n)
%timeit np.einsum('ijk->', a)
100000 loops, best of 3: 3.79 us per loop
%timeit np.sum(a)
100000 loops, best of 3: 7.33 us per loop

n = 100; a = np.arange(n**3, dtype=np.double).reshape(n, n, n)
%timeit np.einsum('ijk->', a)
1000 loops, best of 3: 1.2 ms per loop
%timeit np.sum(a)
1000 loops, best of 3: 1.23 ms per loop

Таким образом, в основном постоянные накладные расходы, а не более быстрый запуск, когда они спускаются к нему.

Ответ 4

Обновление для numpy 1.16.4: Numpy native-функции работают быстрее, чем einsums почти во всех случаях. Только внешний вариант einsum и тест sum23 быстрее, чем версии не einsum.

Если вы можете использовать NumPy нативные функции, сделайте это.

(Изображения, созданные с помощью perfplot, моего проекта.)

Код для воспроизведения участков:

import numpy
import perfplot


def setup1(n):
    return numpy.arange(n, dtype=numpy.double)


def setup2(n):
    return numpy.arange(n ** 2, dtype=numpy.double).reshape(n, n)


def setup3(n):
    return numpy.arange(n ** 3, dtype=numpy.double).reshape(n, n, n)

def setup23(n):
    return (
        numpy.arange(n ** 2, dtype=numpy.double).reshape(n, n),
        numpy.arange(n ** 3, dtype=numpy.double).reshape(n, n, n)
    )


def numpy_sum(a):
    return numpy.sum(a)


def einsum_sum(a):
    return numpy.einsum("ijk->", a)


perfplot.save(
    "sum.png",
    setup=setup3,
    kernels=[numpy_sum, einsum_sum],
    n_range=[2 ** k for k in range(10)],
    logx=True,
    logy=True,
    title="sum",
)


def numpy_power(a):
    return a * a * a


def einsum_power(a):
    return numpy.einsum("ijk,ijk,ijk->ijk", a, a, a)


perfplot.save(
    "power.png",
    setup=setup3,
    kernels=[numpy_power, einsum_power],
    n_range=[2 ** k for k in range(9)],
    logx=True,
    logy=True,
)


def numpy_outer(a):
    return numpy.outer(a, a)


def einsum_outer(a):
    return numpy.einsum("i,k->ik", a, a)


perfplot.save(
    "outer.png",
    setup=setup1,
    kernels=[numpy_outer, einsum_outer],
    n_range=[2 ** k for k in range(13)],
    logx=True,
    logy=True,
    title="outer",
)



def dgemm_numpy(a):
    return numpy.dot(a, a)


def dgemm_einsum(a):
    return numpy.einsum("ij,jk", a, a)


def dgemm_einsum_optimize(a):
    return numpy.einsum("ij,jk", a, a, optimize=True)


perfplot.save(
    "dgemm.png",
    setup=setup2,
    kernels=[dgemm_numpy, dgemm_einsum],
    n_range=[2 ** k for k in range(13)],
    logx=True,
    logy=True,
    title="dgemm",
)



def dot_numpy(a):
    return numpy.dot(a, a)


def dot_einsum(a):
    return numpy.einsum("i,i->", a, a)


perfplot.save(
    "dot.png",
    setup=setup1,
    kernels=[dot_numpy, dot_einsum],
    n_range=[2 ** k for k in range(20)],
    logx=True,
    logy=True,
    title="dot",
)

def sum23_numpy(data):
    a, b = data
    return numpy.sum(a * b)


def sum23_einsum(data):
    a, b = data
    return numpy.einsum('ij,oij->', a, b)


perfplot.save(
    "sum23.png",
    setup=setup23,
    kernels=[sum23_numpy, sum23_einsum],
    n_range=[2 ** k for k in range(10)],
    logx=True,
    logy=True,
    title="sum23",
)