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的四核CPU上运行，以及使用OpenBLAS进行多线程处理，我得到以下结果：

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来加快通过vanilla 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.在任何情况下，为一般输入自动执行此操作似乎只是简单......

Answer 2

einsum有一个'2个操作数，ndim = 2'的专门案例。在这种情况下，有3个操作数，总共3个维度。所以它必须使用通用nditer。

在尝试理解如何解析字符串输入时，我编写了一个纯Python einsum模拟器，https://github.com/hpaulj/numpy-einsum/blob/master/einsum_py.py

（简化）einsum和sum-of-products函数是：

def myeinsum(subscripts, *ops, **kwargs):
    # dropin preplacement for np.einsum (more or less)
    <parse subscript strings>
    <prepare op_axes>
    x = sum_of_prod(ops, op_axes, **kwargs)
    return x

def sum_of_prod(ops, op_axes,...):
    ...
    it = np.nditer(ops, flags, op_flags, op_axes)
    it.operands[nop][...] = 0
    it.reset()
    for (x,y,z,w) in it:
        w[...] += x*y*z
    return it.operands[nop]

使用myeinsum('ik,km,im->i',X,C,X,debug=True)

调试(M,K)=(10,5)的输出

{'max_label': 109, 
 'min_label': 105, 
 'nop': 3, 
 'shapes': [(10, 5), (5, 5), (10, 5)], 
 ....}}
 ...
iter labels: [105, 107, 109],'ikm'

op_axes [[0, 1, -1], [-1, 0, 1], [0, -1, 1], [0, -1, -1]]

如果您在sum-of-prod中编写cython这样的函数，您应该接近广义einsum。

使用完整的(M,K)，这个模拟的einsum慢6-7倍。

基于其他答案的一些时间安排：

In [84]: timeit np.dot(X,C)
1 loops, best of 3: 781 ms per loop

In [85]: timeit np.einsum('ik,km->im',X,C)
1 loops, best of 3: 1.28 s per loop

In [86]: timeit np.einsum('im,im->i',A,X)
10 loops, best of 3: 163 ms per loop

此'im,im->i' step is substantially faster than the other. The sum dimension, m is only 20. I suspect einsum`将此视为一种特殊情况。

In [87]: timeit np.einsum('im,im->i',np.dot(X,C),X)
1 loops, best of 3: 950 ms per loop

In [88]: timeit np.einsum('im,im->i',np.einsum('ik,km->im',X,C),X)
1 loops, best of 3: 1.45 s per loop

这些复合计算的时间只是相应部分的总和。

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

2 个答案: