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

Question

让我们从三个数组开始dtype=np.double。使用 numpy 1.7.1 编译icc并链接到 intel 的mkl. 一个带有 numpy 1.6.1 的 AMD cpugcc也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、似乎至少快两倍。主要的例外是 它从 BLAS 库中调用 DGEMM。那么为什么比其他等效的 numpy 函数更快呢？np.sumaxesnp.dotnp.einsum

完整性的 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

---

主要理论来自np.einsum可以使用SSE2的@sebergs 评论，但 numpy 的 ufunc 直到 numpy 1.8 才会出现（请参阅更改日志）。我相信这是正确的答案，但无法确认。通过更改输入数组的 dtype 并观察速度差异以及并非每个人都观察到相同的时序趋势这一事实，可以找到一些有限的证据。

Answer 1

首先，过去在 numpy 列表上对此进行了很多讨论。例如，参见： http: //numpy-discussion.10968.n7.nabble.com/poor-performance-of-sum-with-sub-machine-word-integer-types-td41.html http://numpy-讨论.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

但是如果我们使用 less limited dtype，我们仍然会得到你期望的结果：

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

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

Answer 2

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

icc为了执行测试，创建了一个新的 python 2.7 安装 -在运行 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.003 秒。其余差异应包含在此问题的其他答案中。

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

所以你在调用时基本上有一个几乎恒定的 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

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

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