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

【问题标题】：Why is numpy's einsum faster than numpy's built in functions?为什么 numpy 的 einsum 比 numpy 的内置函数快？
【发布时间】：2013-08-24 06:38:00
【问题描述】：

让我们从三个dtype=np.double 数组开始。使用icc 编译并链接到英特尔mkl 的numpy 1.7.1 在英特尔CPU 上执行计时。使用gcc 而没有mkl 编译的带有numpy 1.6.1 的AMD cpu 也用于验证时序。请注意，时间与系统大小几乎呈线性关系，并不是由于 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

无论axes 选择如何，np.inner、np.outer、np.kron 和 np.sum 的 Einsum 似乎至少快一倍。主要例外是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 的 ufunc 直到 numpy 1.8 才会使用（参见 change log）。我相信这是正确的答案，但无法确认。通过更改输入数组的 dtype 并观察速度差异以及并非每个人都观察到相同的时序趋势这一事实，可以找到一些有限的证据。

【问题讨论】：

numpy 链接的 BLAS 库是什么？是多线程的吗？
带 AVX 的多线程 MKL BLAS。
顺便提一下，好问题，好例子！这可能值得在邮件列表中询问。之前已经介绍过（尤其是关于sum），但我很惊讶einsum 始终比outer、inner、kron 等快约2 倍。知道在哪里会很有趣差异来自于。
@JoeKington 如果其他人可以重现 ~2x 加速，我想我会将其发布在邮件列表中。奇怪的是，杰米的回答确实证明了这一点。
有点相关：stackoverflow.com/questions/17527340/… 但在这种情况下，速度差异的原因似乎是内存管理，（至少当你开始制作非常大的东西时）

标签： python arrays performance numpy multidimensional-array

【解决方案1】：

首先，过去在 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

【讨论】：

你有关于sum如何挑选累加器的好链接吗？有趣的是，您的x 数组扩展到1E8 元素np.einsum('i->',x,dtype=np.uint64) 仅比sum 快10%（15ms）。
@Ophion - sum 的文档有一些详细信息。您可以使用 dtype kwarg 到 sum 指定它。如果未指定，并且数组的整数 dtype 的精度低于“默认平台整数”（我认为即使在 32 位平台上通常也是 int64），那么它默认为默认整数。见：docs.scipy.org/doc/numpy/reference/generated/numpy.sum.html
另外，sum 是通过np.add.reduce 实现的，所以可以在这里查看ufuncs 的还原源，如果您对细节感兴趣：github.com/numpy/numpy/blob/master/numpy/core/src/umath/…
如果我理解正确，这些是“苹果对苹果”的比较，因为一切都特别限于dtype=np.double?
我想是的。毕竟，这就是你一开始在做的事情。因此，我提出的观点可能并不那么重要！

【解决方案2】：

现在发布了 numpy 1.8，根据文档，所有 ufunc 都应该使用 SSE2，我想仔细检查一下 Seberg 关于 SSE2 的评论是否有效。

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

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

【讨论】：

精彩的跟进！这是我需要更频繁地开始使用einsum 的另一个原因。顺便说一句，在这种情况下，我认为您真的应该将自己的答案标记为正确。

【解决方案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 而不是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

因此，开销几乎是恒定的，而不是一旦开始就可以更快地运行。

【讨论】：

另外，the documentation 表示einsum 也不会进行自动广播，而是依靠用户来表达广播规则进行操作。所以einsum 可能会跳过很多检查（类型检查、广播等）。
奇怪的是它们在我的机器上有所不同，请查看我的编辑。
一维或多维基本上是一样的。 np.sum 调用 np.add.reduce，这是为 1.7 重做的以接受多个轴。因此，在这两种情况下，几乎可以肯定，迭代是由与 C 等效的 np.nditer 非常相似的调用来处理的。除非您避免使用中间数组来执行 numpy 所做的先乘后加的操作，或者您使用的是多线程库，否则您应该会看到除了设置之外的细微差别，这就是我的时间安排所显示的。跨度>
您可能会看到双精度 (SSE) 的 2 倍加速。因为 sum 是幼稚的（可能不会在 1.8+ 上不确定），而 einsum 是专门为使用 SIMD 指令而编写的，但大多数 ufunc 都不是。
@seberg 你说对了，两个处理器都有 SSE2，所以人们期望单精度是 4 倍，而且确实如此。如果你能写出来，我会接受。

【解决方案4】：

numpy 1.21.2 的更新：在几乎所有情况下，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)],
)


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)],
)


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)],
)


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)],
)


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)],
)


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)],
)

【讨论】：

如果您numpy.einsum("ij,jk", a, a, optimize=True)，请在 GEMM 上注明，其性能将相当。延迟更小有点奇怪，这个函数的逻辑是否转移到了 C 中？同样值得尝试np.einsum('i,i->', ...) 和np.einsum('ij,oij->' 以获得更多的苹果与苹果的比较。
@Daniel 添加了这些。