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

my_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)

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

Answer 1

鉴于循环计数（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

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

1 个答案:

基准