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

Давайте начнем с трех массивов dtype=np.double, Синхронизация выполняется на процессоре Intel с использованием numpy 1.7.1, скомпилированной с icc и связано с Intel mkl, Процессор AMD с numpy 1.6.1, скомпилированный с gcc без mkl был также использован для проверки времени. Обратите внимание, что временные рамки масштабируются почти линейно в зависимости от размера системы и не связаны с небольшими накладными расходами, возникающими при использовании функций numpy. 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, но ufuncs numpy будет работать только до версии 1.8 (см. журнал изменений). Я считаю, что это правильный ответ, но не смог подтвердить это. Некоторое ограниченное доказательство можно найти, изменив dtype входного массива и наблюдая разницу скоростей, а также тот факт, что не все наблюдают одинаковые тенденции во времени.

3 ответа

Решение

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

Для выполнения теста была создана новая установка python 2.7 - numpy 1.7 и 1.8 были скомпилированы с icc использование стандартных опций на ядре 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 с. Оставшаяся разница должна быть отражена в других ответах на этот вопрос.

Во-первых, было много прошлых дискуссий по этому поводу в списке 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 является новым и, по-видимому, пытается лучше разобраться с выравниванием кеша и другими проблемами доступа к памяти, в то время как многие из старых функциональных возможностей сосредоточены на легко переносимой реализации по сравнению с сильно оптимизированной. Я просто спекулирую, хотя там.


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

В дополнение к тому, что @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

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

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

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

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

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

Обновление для numpy 1.16.4: собственные функции Numpy работают быстрее, чем 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",
)
Другие вопросы по тегам