为什么numpy的einsum比numpy的内置函数快？

72

让我们从的三个数组开始dtype=np.double。使用numpy 1.7.1编译icc并链接到intel的numpy 1.7.1在intel CPU上执行计时mkl。一个AMD的CPU与编译numpy的1.6.1gcc不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 它从BLAS库调用DGEMM。那么，为什么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才会使用（请参阅更改日志）。我相信这是正确的答案，但无法确认。通过更改输入数组的dtype并观察速度差异以及并非每个人都观察到相同的时序趋势这一事实，可以找到一些有限的证明。

python  arrays  performance  numpy  multidimensional-array 
— 丹尼尔
source
numpy链接到哪个BLAS库？是多线程的吗？
— ali_m
1
带AVX的多线程MKL BLAS。
— 丹尼尔（Daniel）
顺便提一句，很好的问题，很好的例子！可能值得在邮件列表中询问。它之前被覆盖（特别是有关sum），但我很惊讶，einsum一直高于〜2倍快outer，inner，kron等，这将是有趣的知道区别在哪里从何而来。
— 2013年
@JoeKington我想如果其他人可以重现2倍的加速速度，我会把它发布在邮件列表中。奇怪的是，杰米的答案确实证明了这一点。
— 丹尼尔（Daniel）
有点相关：stackoverflow.com/questions/17527340/…但在那种情况下，速度差异的原因似乎是内存管理，（当您开始制作真正的东西时至少要大一点）
— usethedeathstar 2013年

Answers:

32

首先，在numpy列表上有很多关于此的讨论。例如，请参见： http://numpy-discussion.10968.n7.nabble.com/poor-performance-of-sum-with-sub-machine-word-integer-types-td41.html HTTP：// numpy- forum.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
— 乔·金顿
source
您如何sum选择蓄电池有一个很好的链接？有趣的是，将x数组扩展到1E8元素np.einsum('i->',x,dtype=np.uint64)仅比快10％（15毫秒）sum。
— 丹尼尔（Daniel）
@Ophion-有关文档的sum详细信息。您可以使用dtypekwarg来指定它sum。如果未指定，并且数组的整数dtype的精度低于“默认平台整数”（int64我认为通常在32位平台上也是如此），那么它将默认为默认整数。请参阅：docs.scipy.org/doc/numpy/reference/generation/numpy.sum.html
— Joe Kington，
此外，sum是通过实现的，如果您对详细信息感兴趣，请在此处np.add.reduce查看减少ufunc的来源：github.com/numpy/numpy/blob/master/numpy/core/src/umath / ...
— Joe Kington
1
如果我正确理解的话，这些是“苹果对苹果”的比较，因为所有内容都仅限于此dtype=np.double？
— 丹尼尔（Daniel）
2
我认同。毕竟，这是您首先要做的。因此，我提出的观点可能根本不相关！
— 2013年
22

现在numpy 1.8已发布，根据文档，所有ufuncs都应使用SSE2，我想再次检查一下Seberg关于SSE2的评论是否有效。

为了执行该测试，创建了新的python 2.7安装-icc使用运行Ubuntu的AMD opteron核心上的标准选项编译了numpy 1.7和1.8 。

这是在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.003s。其余的差异应包含在该问题的其他答案中。

— 丹尼尔
source
4
很棒的跟进！这是为什么我需要开始einsum更多使用的另一个原因。顺便说一句，在这种情况下，我认为您应该将自己的答案标记为正确。
— 乔·金顿
19

我认为这些时间可以解释发生了什么：

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时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

因此，这几乎是恒定的开销，而一旦他们付诸实施，就不会更快地运行。

— 海梅
source
1
另外，该文档建议einsum也不执行自动广播，并且依赖于用户来表达操作的广播规则。因此，可能有很多einsum可以跳过的检查（类型检查，广播等）。
— 2013年
奇怪的是，它们在我的机器上是不同的，请查看我的编辑。
— 丹尼尔（Daniel）
1个或多个维度基本上是同一件事。np.sum调用np.add.reduce，然后重做1.7以接受多个轴。因此，几乎可以肯定，np.nditer在两种情况下，都可以通过对C等效的非常相似的调用来处理迭代。除非您避免使用中间数组来执行numpy的“乘-然后-加”操作，否则除非您使用多线程库，否则除设置外，您应该看不到什么区别，这就是我的时序显示。
— Jaime
8
您可能应该会看到双精度（SSE）的2倍加速。由于sum天真（可能不确定是否在1.8+上），虽然einsum是专门为使用SIMD指令而编写的，但大多数ufunc并不是这样。
— seberg
1
@seberg您将其钉牢了，两个处理器都具有SSE2，因此可以期望单精度达到其4倍的速度。如果您可以写下来，我会接受。
— 丹尼尔（Daniel）
1

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",
)
— 尼科·施洛默（NicoSchlömer）
source
如果您numpy.einsum("ij,jk", a, a, optimize=True)的性能与GEMM相同，请注意。延迟更小有些奇怪，此函数的逻辑是否移至C？同样值得尝试的np.einsum('i,i->', ...)还有np.einsum('ij,oij->'更多的苹果与苹果的比较。
— 丹尼尔（Daniel）
@Daniel添加了那些。
— NicoSchlömer19年
By using our site, you acknowledge that you have read and understand our Cookie Policy and Privacy Policy.
Licensed under cc by-sa 3.0 with attribution required.