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

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

您可以两全其美：

def func_dot_einsum(C, X):    Y = X.dot(C)    return np.einsum('ij,ij->i', Y, X)

在我的系统上：

In [7]: %timeit func_dot(C, X)10 loops, best of 3: 31.1 ms per loopIn [8]: %timeit func_einsum(C, X)10 loops, best of 3: 105 ms per loopIn [9]: %timeit func_einsum2(C, X)10 loops, best of 3: 43.5 ms per loopIn [10]: %timeit func_dot_einsum(C, X)10 loops, best of 3: 21 ms per loop

如果可用，请

np.dot

使用BLAS，MKL或您拥有的任何库。因此，

np.dot

几乎可以肯定，对的调用是多线程的。

np.einsum

有自己的循环，因此除了使用SIMD来加快原始C实现的速度外，不使用任何优化。

---

然后是一个多输入的einsum调用，运行速度要慢得多…
einsum的numpy源非常复杂，我还不完全了解。因此，请注意以下内容充其量只是推测性的，但这是我认为正在发生的事情…

当您运行时

np.einsum('ij,ij->i', a, b)

，这样做的好处

np.sum(a*b,axis=1)

是避免了必须实例化具有所有产品的中间阵列，并在其上循环两次。因此，从低层次来看，情况是这样的：

for i in range(I):    out[i] = 0    for j in range(J):        out[i] += a[i, j] * b[i, j]

现在说，您正在执行以下 *** 作：

np.einsum('ij,jk,ik->i', a, b, c)

您可以执行与

np.sum(a[:, :, None] * b[None, :, :] * c[:, None, :], axis=(1, 2))

我认为einsum要做的是运行最后的代码，而不必实例化巨大的中间数组，这肯定会使事情变得更快：

In [29]: a, b, c = np.random.rand(3, 100, 100)In [30]: %timeit np.einsum('ij,jk,ik->i', a, b, c)100 loops, best of 3: 2.41 ms per loopIn [31]: %timeit np.sum(a[:, :, None] * b[None, :, :] * c[:, None, :], axis=(1, 2))100 loops, best of 3: 12.3 ms per loop

但是，如果仔细看，摆脱中间存储可能是一件可怕的事情。这是我认为einsum正在做的底层工作：

for i in range(I):    out[i] = 0    for j in range(J):        for k in range(K): out[i] += a[i, j] * b[j, k] * c[i, k]

但是您正在重复大量 *** 作！如果您改为：

for i in range(I):    out[i] = 0    for j in range(J):        temp = 0        for k in range(K): temp += b[j, k] * c[i, k]        out[i] += a[i, j] * temp

您将

I * J * (K-1)

减少乘法运算（和

I *J

额外的加法运算），并节省大量时间。我的猜测是，einsum不够聪明，无法在此级别优化事物。在源代码中有一个提示，它仅优化具有1个或2个 *** 作数的 *** 作，而不是3个。在任何情况下，为通用输入自动执行此 *** 作似乎都不是一件简单的事情…