numpy - einsum vs naive implementation 运行时执行答案

【问题标题】：numpy - einsum vs naive implementation runtime performanednumpy - einsum vs naive implementation 运行时执行
【发布时间】：2020-03-18 14:37:10
【问题描述】：

我有一个大小为 (N,M) 的二维数组 Y，例如：

N, M = 200, 100
Y = np.random.normal(0,1,(N,M))

对于每个 N，我想计算向量 (M,1) 与其转置的点积，它返回一个 (M,M) 矩阵。一种低效的方法是：

Y = Y[:,:,np.newaxis]
[Y[i,:,:] @ Y[i,:,:].T for i in range(N)]

这很慢：第二行的timeit返回

11.7 ms ± 1.39 ms per loop (mean ± std. dev. of 7 runs, 100 loops each)

我认为更好的方法是使用 einsum numpy 函数 (https://docs.scipy.org/doc/numpy/reference/generated/numpy.einsum.html)：

np.einsum('ijk,imk->ijm', Y, Y, optimize=True)

（这意味着：对于每一行 i，创建一个 (j,k) 矩阵，其元素来自最后一维 m 上的点积）

这两种方法确实返回了完全相同的结果，但是这个新版本的运行时间令人失望（速度只有两倍多一点）

3.82 ms ± 146 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

由于第一种方法非常效率低下，因此人们期望通过使用矢量化 einsum 函数获得更多改进...您对此有解释吗？有没有更好的方法来做这个计算？

【问题讨论】：

没有发生总和减少。不要指望 einsum 有帮助。你可以简单地做：Y*Y[:,None,:,0].
没有太多可能的加速（只有大约 2 倍，例如高效的 Numba 实现）。原因是这种计算完全受限于内存带宽。如果你之后做一些减少（例如总和），很容易有一个数量级或更多的可能。优化归约情况的示例：stackoverflow.com/a/58189944/4045774

标签： python performance numpy dot

【解决方案1】：

In [60]: N, M = 200, 100 
    ...: Y = np.random.normal(0,1,(N,M))                                                                             
In [61]: Y1 = Y[:,:,None]

您的迭代，200 步生成 (100,100) 个数组：

In [62]: timeit [Y1[i,:,:]@Y1[i,:,:].T for i in range(N)]                                                            
18.5 ms ± 784 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

einsum 只是稍微快一点：

In [64]: timeit np.einsum('ijk,imk->ijm', Y1,Y1)                                                                     
14.5 ms ± 114 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

但您可以在完整的“批处理”模式下应用@：

In [65]: timeit Y[:,:,None]@Y[:,None,:]                                                                              
7.63 ms ± 224 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

但正如 Divakar 所说，和轴的大小为 1，因此您可以使用普通广播乘法。这是外积，不是矩阵积。

In [66]: timeit Y[:,:,None]*Y[:,None,:]                                                                              
8.2 ms ± 64.7 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

'vectorizing' 在对一个简单的操作进行多次迭代时会带来很大的收益。对于更复杂的操作上的更少操作，增益不是那么大。

【讨论】：

感谢您的回答！我有哪些选择可以尽快完成？

【解决方案2】：

这是一篇旧文章，但包含许多细节：efficient outer product。

特别是如果您有兴趣添加 numba 依赖项，那可能是您最快的选择。

从原帖更新部分 numba 代码并添加多外积：

import numpy as np
from numba import jit
from numba.typed import List

@jit(nopython=True)
def outer_numba(a, b):
    m = a.shape[0]
    n = b.shape[0]
    result = np.empty((m, n))
    for i in range(m):
        for j in range(n):
            result[i, j] = a[i]*b[j]
    return result

@jit(nopython=True)
def multi_outer_numba(Y):

    all_result = List()
    for k in range(Y.shape[0]):
        y = Y[k]           
        n = y.shape[0]
        tmp_res = np.empty((n, n))
        for i in range(n):
            for j in range(n):
                tmp_res[i, j] = y[i]*y[j]
        all_result.append(tmp_res)

    return all_result

r = [outer_numba(Y[i],Y[i]) for i in range(N)]
r = multi_outer_numba(Y)

【讨论】：