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

Question

我们从三个dtype=np.double的数组开始。时间是在使用numpy 1.7.1编译的icc并链接到英特尔的mkl的英特尔CPU上进行的。还使用了一个带有numpy 1.6.1编译的gcc但没有mkl的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

对于np.inner，np.outer，np.kron和np.sum，无论选择哪个axes，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 1.8（请参见change log），numpy的ufuncs才能使用它。我相信这是正确的答案，但我无法确认。一些有限的证据可以通过改变输入数组的dtype并观察速度差异以及不是每个人都观察到相同的时间趋势来找到。

- Daniel

---

NumPy链接的是哪个BLAS库？它是否支持多线程？ - ali_m

使用AVX的多线程MKL BLAS。 - Daniel

顺便说一句，很好的问题，而且有很好的例子！也许值得在邮件列表上提问。这个问题以前已经讨论过（特别是关于sum），但我对einsum始终比outer、inner、kron等快大约两倍感到惊讶。了解这种差异的来源将会很有趣。 - Joe Kington

@JoeKington 如果有其他人能够复现这个大约2倍的加速，我想我会在邮件列表上发布它。奇怪的是，Jamie的回答确实证明了这一点。 - Daniel

@usethedeathstar 很好的观点；然而，这里绝对不是内存问题。 - Daniel

Answer 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

- Joe Kington

---

你有关于sum如何选择累加器的好链接吗？有趣的是，当你的x数组扩展到1E8个元素时，np.einsum('i->'，x，dtype = np.uint64)只比sum快大约10％（15毫秒）。 - Daniel

@Ophion - sum 的文档有一些细节。您可以使用 sum 的 dtype 关键字参数指定它。如果未指定，并且数组具有比“默认平台整数”（通常在32位平台上为int64）低精度的整数 dtype，则会默认使用默认整数。请参阅：http://docs.scipy.org/doc/numpy/reference/generated/numpy.sum.html - Joe Kington

另外，sum是通过np.add.reduce实现的，如果您对细节感兴趣，请查看此处有关缩减ufunc的源代码：https://github.com/numpy/numpy/blob/master/numpy/core/src/umath/reduction.c - Joe Kington

如果我理解正确的话，这些都是“苹果对苹果”的比较，因为所有内容都被明确限制为 dtype=np.double？ - Daniel

我想是这样的。毕竟，这本来就是你一开始要做的事情。因此，我提出的观点可能并不那么相关！ - Joe Kington

Answer 2

现在numpy 1.8已经发布，根据文档，所有ufunc应该使用SSE2，我想要再次确认Seberg有关SSE2的评论是否有效。

为了进行测试，创建了一个新的python 2.7安装程序-使用标准选项在AMD Opteron核心运行Ubuntu上编译了numpy 1.7和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秒。其余的差异应在本问题的其他答案中得到解决。

- Daniel

---

跟进得非常好！这是我需要更经常开始使用 einsum 的又一个原因。顺便说一句，我认为在这种情况下你应该将自己的答案标记为正确的。 - Joe Kington

Answer 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时，基本上你几乎有一个近乎恒定的3微秒的开销，因此它们基本上运行得一样快，但其中一个需要更长时间才能启动。这可能是为什么？我猜测原因如下：

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

因此，大多数情况下需要付出固定的开销，而不是一旦开始就可以更快地运行。

- Jaime

---

此外，文档 还建议 einsum 不会执行自动广播，而是依赖用户来表达操作的广播规则。因此，einsum 可能会跳过许多检查（类型检查、广播等）。 - ely

一个或多个维度基本上是相同的事情。np.sum 调用 np.add.reduce，在 1.7 版本中对其进行了重新设计以接受多个轴。因此，在这两种情况下，迭代几乎肯定是由非常类似于 np.nditer 的 C 等效调用处理的。除非您避免使用中间数组来执行 numpy 执行的乘-加操作，或者您正在使用多线程库，否则除了设置之外，您应该看不到太大的差异，这也是我的时间测试所显示的。 - Jaime

顺便说一下，我可以在 AMD CPU 上使用 einsum 实现 2 倍加速，而不需要链接 mkl 和 numpy 1.6.1。还有其他人看到和我一样的时间趋势吗？ - Daniel

双精度（SSE）可能会使运行速度加快两倍。因为sum函数比较原始（不确定是否在1.8以上版本），而einsum函数则专门使用SIMD指令编写，大多数ufuncs不使用这些指令。 - seberg

@seberg，你说得对，两个处理器都有SSE2指令集，因此单精度运算的速度应该是四倍于双精度的。如果你能把这个写出来，我会接受它。 - Daniel

Answer 4

numpy 1.21.2的更新：在几乎所有情况下，Numpy的本地函数比einsums更快。只有einsum的outer变体和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)],
)

- Nico Schlömer

---

关于GEMM的一点说明，如果您使用numpy.einsum("ij,jk", a, a, optimize=True)，性能将是等效的。有些奇怪的是延迟更小了，这个函数的逻辑是否移动到了C中？同时也值得尝试np.einsum('i,i->', ...)以及np.einsum('ij,oij->'进行更为公正的比较。 - Daniel

@Daniel 已添加了这些。 - Nico Schlömer