Python 优化现有的3D Numpy矩阵乘法

我有一些刚刚完成的代码。它按预期工作。我选择在Numpy中使用

点

，因为根据我有限的经验，如果系统上安装了BLAS，它比通常形式的矩阵乘法更快。然而，你会注意到，我不得不转置很多东西。我确信这实际上超过了使用

dot

的好处

我提供了论文中发现的数学方程。请注意，这是一个递归

以下是我的代码实现：

我首先提供必要的部件及其尺寸

P = array([[0.73105858, 0.26894142],
           [0.26894142, 0.73105858]])  # shape (K,K)

B = array([[6.07061629e-09, 0.00000000e+00],
           [0.00000000e+00, 2.57640371e-10]])  # shape (K,K)

dP = array([[[ 0.19661193, -0.19661193],
             [ 0.        ,  0.        ]],

           [[ 0.        ,  0.        ],
            [ 0.19661193, -0.19661193]],

           [[ 0.        ,  0.        ],
            [ 0.        ,  0.        ]],

           [[ 0.        ,  0.        ],
            [ 0.        ,  0.        ]]])  # shape (L,K,K)

dB = array([[[ 0.00000000e+00,  0.00000000e+00],
             [ 0.00000000e+00,  0.00000000e+00]],

            [[ 0.00000000e+00,  0.00000000e+00],
             [ 0.00000000e+00,  0.00000000e+00]],

            [[-1.16721049e-09,  0.00000000e+00],
             [ 0.00000000e+00,  0.00000000e+00]],

            [[ 0.00000000e+00,  0.00000000e+00],
             [ 0.00000000e+00, -1.27824683e-09]]])  # shape (L,K,K)

a = array([[[ 7.60485178e-08,  5.73923956e-07]],

           [[-5.54100398e-09, -8.75213012e-08]],

           [[-1.25878643e-08, -1.48361081e-07]],

           [[-2.73494035e-08, -1.74585971e-07]]])  # shape (L,1,K)

alpha = array([[0.11594542, 0.88405458]])  # shape (1,K)

c = 1  # scalar

这是实际的Numpy计算。注意所有转置的用法

term1 = (a/c).dot(P).dot(B)
term2 = dP.transpose((0,2,1)).dot(alpha.T).transpose((0,2,1)).dot(B)
term3 = dB.dot(  P.T.dot(alpha.T) ).transpose((0,2,1))
a = term1 + term2 + term3

那么你应该得到：

>>> a
array([[[ 1.38388584e-10, -5.87312190e-12]],

       [[ 1.05516813e-09, -4.47819530e-11]],

       [[-3.76451117e-10, -2.88160549e-17]],

       [[-4.06412069e-16, -8.65984406e-10]]])

请注意，

alpha

和

a

的形状是我选择的。如果您发现它可以提供优异的性能，则可以对其进行更改

我想指出的是，我认为现有的代码是快速的。实际上，非常快。然而，我一直想知道是否有人能做得更好。请试一下。我已经分析了我的代码（它有很多Numpy广播和矢量化），这不一定是一个瓶颈，因为在我非常旧的机器上评估需要23微秒。然而，这只是递归的一个步骤。这意味着它将按顺序计算

N次

。因此，即使是最小的增益也会对大序列产生很大的影响

谢谢你抽出时间

编辑/更新： 感谢@max9111，他建议我看看这个问题，我已经管理了一些比

a

的Numpy计算更快的Numba代码。它需要14微秒，而不是原来的23微秒

这是：

import numba as nb
@nb.njit(fastmath=True,parallel=True,boundscheck=False)
def get_next_a(a,alpha,P,dP,B,dB,c):
    N,M,_ = dP.shape
    new_a = np.zeros((N,1,M),dtype=np.float64)
    new_a = np.zeros((N,1,M))
    entry = 0
    for idx in nb.prange(N):
        for i in range(M):
            for j in range(M):
                term1 =  a[idx,0,j]*P[j,i]*B[i,i]/c
                term2 = alpha[0,j]*dP[idx,j,i]*B[i,i] 
                term3 = alpha[0,j]*P[j,i]*dB[idx,i,i]
                entry += term1 + term2 + term3
            new_a[idx,0,i] = entry
            entry = 0
    return new_a

您将看到

get\u next\u a

返回所需的结果。然而，当我在包含Numpy的纯python函数中调用它时，它会抱怨。以下是我的实际代码片段：

def forward_recursions(X,working_params):

#    P,dP,B,dB,pi,dpi = setup(X,working_params) 
    # Dummy Data and Parameters instead of setup
    working_params = np.random.uniform(0,2,size=100)
    X = np.random.uniform(0,1,size=1000)
    P = np.random.uniform(0,1,size=(10,10))
    norm = P.sum(axis=1)
    P = P/norm[:,None]
    dP = np.random.uniform(-1,1,size=(100,10,10))
    # We pretend that all 1000 of the 10 x 10 matrices 
    # only have diagonal entries
    B = np.random.uniform(0,1,size=(1000,10,10)) 
    dB = np.random.uniform(0,1,size=(1000,100,10,10))
    pi = np.random.uniform(0,1,size=10)
    norm = pi.sum()
    pi = (pi/norm).reshape(1,10)
    dpi = np.random.uniform(0,1,size=(100,1,10))

    T = len(X)
    N = len(working_params)
    M = np.int(np.sqrt(N))
    ones = np.ones((M,1))


    alpha = pi.dot(B[0])
    scale = alpha.dot(ones)
    alpha = alpha/scale
    ll = np.log(scale)
    a = dpi.dot(B[0]) + dB[0].dot(pi.T).transpose((0,2,1))
    for t in range(1,T):

        old_scale = scale
        alpha = alpha.dot(P).dot(B[t])
        scale = alpha.dot(ones)
        ll += np.log(scale)
        alpha = alpha/scale

        # HERE IS THE NUMBA FUNCTION

        a = get_next_a(a,alpha,P,dP,B[t],dB[t],old_scale)

    dll = a.dot(ones).reshape((N,1))/scale
    return ll,dll,a

我知道包含我自己的代码取决于未包含的其他函数，因此意味着

前向递归将不会运行。我只是希望它能给我们一些观点

我得到的错误是

TypingError: Invalid use of Function(<built-in function iadd>) with argument(s) of type(s): (Literal[int](0), array(float64, 2d, C))
Known signatures:
 * (int64, int64) -> int64
 * (int64, uint64) -> int64
 * (uint64, int64) -> int64
 * (uint64, uint64) -> uint64
 * (float32, float32) -> float32
 * (float64, float64) -> float64
 * (complex64, complex64) -> complex64
 * (complex128, complex128) -> complex128
 * parameterized
In definition 0:
    All templates rejected with literals.
In definition 1:
    All templates rejected without literals.
In definition 2:
    All templates rejected with literals.
In definition 3:
    All templates rejected without literals.
In definition 4:
    All templates rejected with literals.
In definition 5:
    All templates rejected without literals.
In definition 6:
    All templates rejected with literals.
In definition 7:
    All templates rejected without literals.
In definition 8:
    All templates rejected with literals.
In definition 9:
    All templates rejected without literals.
In definition 10:
    All templates rejected with literals.
In definition 11:
    All templates rejected without literals.
In definition 12:
    All templates rejected with literals.
In definition 13:
    All templates rejected without literals.
In definition 14:
    All templates rejected with literals.
In definition 15:
    All templates rejected without literals.
This error is usually caused by passing an argument of a type that is unsupported by the named function.
[1] During: typing of intrinsic-call at <ipython-input-251-50e636317ef8> (13)

TypingError:函数（）与类型为（Literal[int]（0），数组（float64，2d，C）的参数的使用无效
已知签名：
*（int64，int64）->int64
*（int64，uint64）->int64
*（uint64，int64）->int64
*（uint64，uint64）->uint64
*（float32，float32）->float32
*（float64，float64）->float64
*（complex64，complex64）->complex64
*（complex128，complex128）->complex128
*参数化
在定义0中：
所有模板均拒绝使用文字。
在定义1中：
拒绝所有没有文字的模板。
在定义2中：
所有模板均拒绝使用文字。
在定义3中：
拒绝所有没有文字的模板。
在定义4中：
所有模板均拒绝使用文字。
在定义5中：
拒绝所有没有文字的模板。
在定义6中：
所有模板均拒绝使用文字。
在定义7中：
拒绝所有没有文字的模板。
在定义8中：
所有模板均拒绝使用文字。
在定义9中：
拒绝所有没有文字的模板。
在定义10中：
所有模板均拒绝使用文字。
在定义11中：
拒绝所有没有文字的模板。
在定义12中：
所有模板均拒绝使用文字。
在定义13中：
拒绝所有没有文字的模板。
在定义14中：
所有模板均拒绝使用文字。
在定义15中：
拒绝所有没有文字的模板。
此错误通常由传递指定函数不支持的类型的参数引起。
[1] 期间：在（13）处键入内部调用

我不明白这是什么意思。你知道我怎么才能解决这样的问题吗。非常感谢您抽出时间

Q：…如果可以做得更好

您的原样代码在我（似乎更老）的机器上执行，不是在发布的
~23[us]
中，而是在第一次调用时执行
~45[ms]
，并且享受
iCACHE
和
dCACHE
层次结构中
~77..1x[us]
之间的所有预取：

>>> from zmq import Stopwatch; aClk = Stopwatch()
>>> import numpy as np
>>>
>>> aClk.start(); a = ( a / c ).dot( P ).dot( B ) + dP.transpose( ( 0, 2, 1) ).dot( alpha.T ).transpose( ( 0, 2, 1 ) ).dot( B ) + dB.dot(  P.T.dot( alpha.T ) ).transpose( ( 0, 2, 1 ) ); aClk.stop()
44679
>>> aClk.start(); a = ( a / c ).dot( P ).dot( B ) + dP.transpose( ( 0, 2, 1) ).dot( alpha.T ).transpose( ( 0, 2, 1 ) ).dot( B ) + dB.dot(  P.T.dot( alpha.T ) ).transpose( ( 0, 2, 1 ) ); aClk.stop()
149
>>> aClk.start(); a = ( a / c ).dot( P ).dot( B ) + dP.transpose( ( 0, 2, 1) ).dot( alpha.T ).transpose( ( 0, 2, 1 ) ).dot( B ) + dB.dot(  P.T.dot( alpha.T ) ).transpose( ( 0, 2, 1 ) ); aClk.stop()
113
>>> aClk.start(); a = ( a / c ).dot( P ).dot( B ) + dP.transpose( ( 0, 2, 1) ).dot( alpha.T ).transpose( ( 0, 2, 1 ) ).dot( B ) + dB.dot(  P.T.dot( alpha.T ) ).transpose( ( 0, 2, 1 ) ); aClk.stop()
128
>>> aClk.start(); a = ( a / c ).dot( P ).dot( B ) + dP.transpose( ( 0, 2, 1) ).dot( alpha.T ).transpose( ( 0, 2, 1 ) ).dot( B ) + dB.dot(  P.T.dot( alpha.T ) ).transpose( ( 0, 2, 1 ) ); aClk.stop()
82
>>> aClk.start(); a = ( a / c ).dot( P ).dot( B ) + dP.transpose( ( 0, 2, 1) ).dot( alpha.T ).transpose( ( 0, 2, 1 ) ).dot( B ) + dB.dot(  P.T.dot( alpha.T ) ).transpose( ( 0, 2, 1 ) ); aClk.stop()
100
>>> aClk.start(); a = ( a / c ).dot( P ).dot( B ) + dP.transpose( ( 0, 2, 1) ).dot( alpha.T ).transpose( ( 0, 2, 1 ) ).dot( B ) + dB.dot(  P.T.dot( alpha.T ) ).transpose( ( 0, 2, 1 ) ); aClk.stop()
77
>>> aClk.start(); a = ( a / c ).dot( P ).dot( B ) + dP.transpose( ( 0, 2, 1) ).dot( alpha.T ).transpose( ( 0, 2, 1 ) ).dot( B ) + dB.dot(  P.T.dot( alpha.T ) ).transpose( ( 0, 2, 1 ) ); aClk.stop()
97
>>> a
array([[[ 1.38387304e-10, -5.87323502e-12]],
       [[ 1.05516829e-09, -4.47819355e-11]],
       [[-3.76450816e-10, -2.60843400e-20]],
       [[-1.41384088e-18, -8.65984377e-10]]])

>>> aClk.start(); a = ( a / c ).dot( P ).dot( B ) + dP.transpose( ( 0, 2, 1) ).dot( alpha.T ).transpose( ( 0, 2, 1 ) ).dot( B ) + dB.dot(  P.T.dot( alpha.T ) ).transpose( ( 0, 2, 1 ) ); aClk.stop()
97
>>> a
array([[[ 1.38387304e-10, -5.87323502e-12]],    
       [[ 1.05516829e-09, -4.47819355e-11]],
       [[-3.76450816e-10, -2.60843400e-20]],
       [[-1.41384088e-18, -8.65984377e-10]]])

有趣的是，多次重新运行代码，将处理结果重新分配回
a实际上不会改变
a中的值：

>>> from zmq import Stopwatch; aClk = Stopwatch()
>>> import numpy as np
>>>
>>> aClk.start(); a = ( a / c ).dot( P ).dot( B ) + dP.transpose( ( 0, 2, 1) ).dot( alpha.T ).transpose( ( 0, 2, 1 ) ).dot( B ) + dB.dot(  P.T.dot( alpha.T ) ).transpose( ( 0, 2, 1 ) ); aClk.stop()
44679
>>> aClk.start(); a = ( a / c ).dot( P ).dot( B ) + dP.transpose( ( 0, 2, 1) ).dot( alpha.T ).transpose( ( 0, 2, 1 ) ).dot( B ) + dB.dot(  P.T.dot( alpha.T ) ).transpose( ( 0, 2, 1 ) ); aClk.stop()
149
>>> aClk.start(); a = ( a / c ).dot( P ).dot( B ) + dP.transpose( ( 0, 2, 1) ).dot( alpha.T ).transpose( ( 0, 2, 1 ) ).dot( B ) + dB.dot(  P.T.dot( alpha.T ) ).transpose( ( 0, 2, 1 ) ); aClk.stop()
113
>>> aClk.start(); a = ( a / c ).dot( P ).dot( B ) + dP.transpose( ( 0, 2, 1) ).dot( alpha.T ).transpose( ( 0, 2, 1 ) ).dot( B ) + dB.dot(  P.T.dot( alpha.T ) ).transpose( ( 0, 2, 1 ) ); aClk.stop()
128
>>> aClk.start(); a = ( a / c ).dot( P ).dot( B ) + dP.transpose( ( 0, 2, 1) ).dot( alpha.T ).transpose( ( 0, 2, 1 ) ).dot( B ) + dB.dot(  P.T.dot( alpha.T ) ).transpose( ( 0, 2, 1 ) ); aClk.stop()
82
>>> aClk.start(); a = ( a / c ).dot( P ).dot( B ) + dP.transpose( ( 0, 2, 1) ).dot( alpha.T ).transpose( ( 0, 2, 1 ) ).dot( B ) + dB.dot(  P.T.dot( alpha.T ) ).transpose( ( 0, 2, 1 ) ); aClk.stop()
100
>>> aClk.start(); a = ( a / c ).dot( P ).dot( B ) + dP.transpose( ( 0, 2, 1) ).dot( alpha.T ).transpose( ( 0, 2, 1 ) ).dot( B ) + dB.dot(  P.T.dot( alpha.T ) ).transpose( ( 0, 2, 1 ) ); aClk.stop()
77
>>> aClk.start(); a = ( a / c ).dot( P ).dot( B ) + dP.transpose( ( 0, 2, 1) ).dot( alpha.T ).transpose( ( 0, 2, 1 ) ).dot( B ) + dB.dot(  P.T.dot( alpha.T ) ).transpose( ( 0, 2, 1 ) ); aClk.stop()
97
>>> a
array([[[ 1.38387304e-10, -5.87323502e-12]],
       [[ 1.05516829e-09, -4.47819355e-11]],
       [[-3.76450816e-10, -2.60843400e-20]],
       [[-1.41384088e-18, -8.65984377e-10]]])

>>> aClk.start(); a = ( a / c ).dot( P ).dot( B ) + dP.transpose( ( 0, 2, 1) ).dot( alpha.T ).transpose( ( 0, 2, 1 ) ).dot( B ) + dB.dot(  P.T.dot( alpha.T ) ).transpose( ( 0, 2, 1 ) ); aClk.stop()
97
>>> a
array([[[ 1.38387304e-10, -5.87323502e-12]],    
       [[ 1.05516829e-09, -4.47819355e-11]],
       [[-3.76450816e-10, -2.60843400e-20]],
       [[-1.41384088e-18, -8.65984377e-10]]])

这意味着，代码按原样做了大量工作，以便最终提供一个不变值
a（一个重新生成的身份，代价是花费~XY[us]这样做-您是唯一一个决定是否适合您的目标应用程序的人）


关于希望有改进空间的客户的评论：
好吧，考虑到
N
~1E（3..6）
和
K~10
和
L~100
，对任何改进工作都没有太大的期望，这些改进工作是为了重新解决（到目前为止，
的身份验证结果）希望提高性能的问题

寻求改进的目标处理将按顺序重复超过
~1000x
…小于
~1000000x
，这意味着：


RAM绑定问题不是主要问题，因为静态部分上的缓存效应（所有静态部分的大小只有少数
[MB]
）将以尽可能短的延迟从缓存中重新使用
在
numpy
-工具的设计和工程设计中已经预先解决了与CPU有关的问题（在可行的情况下，利用CPU的SIMD资源和向量对齐跨步技巧-因此对用户级编码的期望不高）

最后，但并非最不重要的一点是，人们可能会评论转置的“成本””——
numpy除了改变索引顺序之外，对转置矩阵没有其他作用-
>>> %timeit forward_recursions_numba(X,P,dP,B,dB,pi,dpi)
51.3 ms ± 389 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

>>> %timeit forward_recursions(X,P,dP,B,dB,pi,dpi)
271 ms ± 1.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)