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

Question

通常，我从numpy的einsum函数获得了良好的性能（我喜欢它的语法）。@Ophion对这个问题的回答表明-在测试的情况下-
einsum始终优于“内置”功能（有时会稍微好一些，有时会很多）。但是我刚遇到einsum慢得多的情况。考虑以下等效功能：

(M, K) = (1000000, 20)
C = np.random.rand(K, K)
X = np.random.rand(M, K)

def func_dot(C, X):
    Y = X.dot(C)
    return np.sum(Y * X, axis=1)

def func_einsum(C, X):
    return np.einsum('ik,km,im->i', X, C, X)

def func_einsum2(C, X):
    # Like func_einsum but break it into two steps.
    A = np.einsum('ik,km', X, C)
    return np.einsum('ik,ik->i', A, X)

我期望func_einsum跑得最快，但这不是我遇到的。在具有超线程，numpy版本1.9.0.dev-7ae0206和带有OpenBLAS的多线程的四核CPU上运行，我得到以下结果：

In [2]: %time y1 = func_dot(C, X)
CPU times: user 320 ms, sys: 312 ms, total: 632 ms
Wall time: 209 ms
In [3]: %time y2 = func_einsum(C, X)
CPU times: user 844 ms, sys: 0 ns, total: 844 ms
Wall time: 842 ms
In [4]: %time y3 = func_einsum2(C, X)
CPU times: user 292 ms, sys: 44 ms, total: 336 ms
Wall time: 334 ms

当我增加到K200时，差异更加极端：

In [2]: %time y1= func_dot(C, X)
CPU times: user 4.5 s, sys: 1.02 s, total: 5.52 s
Wall time: 2.3 s
In [3]: %time y2= func_einsum(C, X)
CPU times: user 1min 16s, sys: 44 ms, total: 1min 16s
Wall time: 1min 16s
In [4]: %time y3 = func_einsum2(C, X)
CPU times: user 15.3 s, sys: 312 ms, total: 15.6 s
Wall time: 15.6 s

有人可以解释为什么einsum这么慢吗？

如果重要的话，这是我的numpy配置：

In [6]: np.show_config()
lapack_info:
    libraries = ['openblas']
    library_dirs = ['/usr/local/lib']
    language = f77
atlas_threads_info:
    libraries = ['openblas']
    library_dirs = ['/usr/local/lib']
    define_macros = [('ATLAS_WITHOUT_LAPACK', None)]
    language = c
    include_dirs = ['/usr/local/include']
blas_opt_info:
    libraries = ['openblas']
    library_dirs = ['/usr/local/lib']
    define_macros = [('ATLAS_INFO', '"\\"None\\""')]
    language = c
    include_dirs = ['/usr/local/include']
atlas_blas_threads_info:
    libraries = ['openblas']
    library_dirs = ['/usr/local/lib']
    define_macros = [('ATLAS_INFO', '"\\"None\\""')]
    language = c
    include_dirs = ['/usr/local/include']
lapack_opt_info:
    libraries = ['openblas', 'openblas']
    library_dirs = ['/usr/local/lib']
    define_macros = [('ATLAS_WITHOUT_LAPACK', None)]
    language = f77
    include_dirs = ['/usr/local/include']
lapack_mkl_info:
  NOT AVAILABLE
blas_mkl_info:
  NOT AVAILABLE
mkl_info:
  NOT AVAILABLE

Answer 1

您可以两全其美：

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 loop

In [8]: %timeit func_einsum(C, X)
10 loops, best of 3: 105 ms per loop

In [9]: %timeit func_einsum2(C, X)
10 loops, best of 3: 43.5 ms per loop

In [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 loop

In [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个。在任何情况下，为通用输入自动执行此操作似乎都不是一件简单的事情…