Python 双np.einsum的性能及如何加速

考虑这一点：

我的_func（）是我第一次尝试进行计算，但我想加快计算速度。然后，我使用以下修改过的函数进行了尝试：

def my_func_2(a,b,c):    
    OuterSum = np.einsum('lpk, lkm, lp -> lm', a, b, c)
    Result = 2 * OuterSum
    return Result

但是，当我在两个函数上运行

%timeit

时，我得到

%timeit my_func(a,b,c)
293 µs ± 1.37 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit my_func_2(a,b,c)
347 µs ± 1.47 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

为什么第二种方法比第一种慢？如何优化我的_func（）使其更快？

鉴于循环计数（a和b的长度）与沿其他轴的长度相比不是一个很大的数字，我们可以运行一个简单的循环，并在每次迭代中利用BLAS支持的矩阵乘法。长度还意味着每次迭代都有足够的总和缩减，这证明了for循环在这种情况下的合理性

执行工作将是：-

N,M = b.shape[::2]
out = np.empty((N,M))
for i in range(N):
    out[i] = c[i].dot(a[i]).dot(b[i])
out *= 2

标杆管理 使用

optimize

参数，这似乎显著提高了

my_func_2

的性能，并将建议的参数作为另一个函数添加进来-

def my_func(a,b,c, optimize=False):
    InnerSum = np.einsum('lpk, lkm -> lpm', a, b,optimize=optimize)
    OuterSum = np.einsum('lp, lpm -> lm', c, InnerSum, optimize=optimize)
    Result = 2 * OuterSum
    return Result

def my_func_2(a,b,c, optimize=False):    
    OuterSum = np.einsum('lpk, lkm, lp -> lm', a, b, c,optimize=optimize)
    Result = 2 * OuterSum
    return Result

def my_func_3(a,b,c):
    N,M = b.shape[::2]
    out = np.empty((N,M))
    for i in range(N):
        out[i] = c[i].dot(a[i]).dot(b[i])
    out *= 2
    return out

时间安排-

In [51]: # Setup used in the question
    ...: np.random.seed(0)
    ...: a = np.random.uniform(0,1,size=[14,25,25])
    ...: b = np.random.uniform(0,1,size=[14,25,25])
    ...: c = np.random.uniform(0,1,size=[14,25])

# With einsum optimize set as False
In [52]: %timeit my_func(a,b,c, optimize=False)
    ...: %timeit my_func_2(a,b,c, optimize=False)
    ...: %timeit my_func_3(a,b,c)
1000 loops, best of 3: 255 µs per loop
1000 loops, best of 3: 302 µs per loop
10000 loops, best of 3: 28.7 µs per loop

# With einsum optimize set as True
In [53]: %timeit my_func(a,b,c, optimize=True)
    ...: %timeit my_func_2(a,b,c, optimize=True)
    ...: %timeit my_func_3(a,b,c)
1000 loops, best of 3: 334 µs per loop
10000 loops, best of 3: 77.6 µs per loop
10000 loops, best of 3: 28.6 µs per loop

---

尝试将

optimize

设置为

True

：

np.einsum（'lpk，lkm，lp->lm'，a，b，c，optimize=True）

？在我的笔记本电脑上

my_func_2

大约快10倍。无论如何，看看这个项目，了解一下

optimize=True

在

np.einsum

中做了什么：

optimize=True

将是mm；然而，像在

函数3

中那样构建循环GEMM是不够聪明的

my_func

with

optimize=True

只会增加约100 us的额外开销，如图所示。@Daniel你对func_3不够聪明到底是什么意思？同意我的函数部分。

optimize=True

只能处理

tensordot

和规范

einsum

之间的选择。这里的最佳解决方案是引入一个基于python的for循环，

optimize=True

不会覆盖它。@Daniel不认为一个带有

optimize=True

的循环会比基于点的循环做得更好，如果这是你的意思的话。不，除非形状使

（a*b）*c

vs

a*（b*c）

一种明显更好的收缩方式。只是想指出为什么

optimize=True

在

myfunc

中更贵，而在

myfunc2

中更便宜。

In [51]: # Setup used in the question
    ...: np.random.seed(0)
    ...: a = np.random.uniform(0,1,size=[14,25,25])
    ...: b = np.random.uniform(0,1,size=[14,25,25])
    ...: c = np.random.uniform(0,1,size=[14,25])

# With einsum optimize set as False
In [52]: %timeit my_func(a,b,c, optimize=False)
    ...: %timeit my_func_2(a,b,c, optimize=False)
    ...: %timeit my_func_3(a,b,c)
1000 loops, best of 3: 255 µs per loop
1000 loops, best of 3: 302 µs per loop
10000 loops, best of 3: 28.7 µs per loop

# With einsum optimize set as True
In [53]: %timeit my_func(a,b,c, optimize=True)
    ...: %timeit my_func_2(a,b,c, optimize=True)
    ...: %timeit my_func_3(a,b,c)
1000 loops, best of 3: 334 µs per loop
10000 loops, best of 3: 77.6 µs per loop
10000 loops, best of 3: 28.6 µs per loop