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

Question

我通常从 numpy 的 einsum 函数中获得良好的性能(我喜欢它的语法)。 @Ophion 对 this question 的回答表明 - 对于测试的案例 - 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

当我将 K 增加到 200 时，差异更加极端:

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 不够聪明，无法在此级别上优化事物。在source code有迹象表明它只优化了 1 或 2 个操作数的操作，而不是 3 个。在任何情况下，为一般输入自动执行此操作似乎并不简单......