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

Question

让我们从三个 dtype=np.double 数组开始。使用用 icc 编译并链接到 intel 的 mkl 的 numpy 1.7.1 在 intel 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

np.inner、np.outer、np.kron 和 的

Einsum 似乎至少快两倍np.sum 与 axes 选择无关。主要异常(exception)是 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 并观察速度差异以及并非每个人都观察到相同的时序趋势这一事实，可以找到一些有限的证据。

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