双人表演np.einsum公司以及如何加速

Question

考虑一下这个MWE：

import numpy as np
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])

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

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

^{pr2}$

但是，当我对这两个函数运行%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)

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

Answer 1

考虑到循环计数（a和{}的长度）与沿其他轴的长度相比不是一个很大的数字，我们可以运行一个简单的循环，并在每次迭代中利用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的性能，并且还将所建议的一个作为另一个函数添加进来-

^{pr2}$

时间安排-

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