时间戳

首页 > 解决方案 > 优化现有的 3D Numpy 矩阵乘法

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

问题描述

我有一些刚刚完成的代码。它按预期工作。我选择dot在 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]]])

请注意，我已经选择了alphaas的形状。a如果您发现它可以提供卓越的性能，则可以更改这些。

我想指出，我认为现有的代码很快。实际上，非常快。但是，我一直想知道是否可以做得更好。请试一试。我已经分析了我的代码（其中有很多 Numpy 广播和矢量化），这不一定是瓶颈，因为在我非常旧的机器上评估需要 23 微秒。但是，它是递归的一个步骤。这意味着它N以顺序方式进行评估。因此，即使是最小的增益也会对大序列产生很大的影响。

感谢您的时间。

编辑/更新：

感谢@max9111，他建议我在这里查看这个问题，我已经管理了一些运行速度比 Numpy 计算更快的 Numba 代码a。它需要 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_next_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

我知道包含我自己的代码取决于未包含的其他功能，因此意味着forward_recursions不会运行。我只是希望它能给一些观点。

我得到的错误是

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)

我不明白这是什么意思。你也许知道我怎么能解决这样的问题。非常感谢您的宝贵时间。

标签： pythonnumpymultidimensional-arraymatrix-multiplicationtensor

解决方案

问：……如果可以做得更好？

您的原样代码在我的（似乎更老的）机器上执行，而不是在张贴的机器上~ 23 [us]，而是~ 45 [ms]在第一次调用时执行，并且享受所有预取iCACHE和dCACHE介于两者之间的层次结构~77..1xx [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]]])

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

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

这意味着，代码按原样做了很多工作，以便最终提供一个不变的值( 重新生成的身份，代价是花费 ~ XY [我们] 这样做 - 你是唯一一个决定，你的目标应用程序是否可以）a

关于希望有改进空间的评论：

好吧，鉴于and ，对任何改进工作都没有太多期望，赞助重新解决（到目前为止 an的身份结果）希望提高性能。N ~ 1E(3..6)K ~ 10L ~ 100a

寻求改进的目标处理将按顺序重复多于~1,000x……少于~ 1,000,000x，这意味着：

  • RAM 绑定问题不是主要问题，因为缓存对静态部分的影响，所有大小都只有几个[MB]，将以尽可能短的延迟从缓存中重用
  • CPU 绑定的问题已经在 -tools 的设计和工程中预先解决numpy（在可行的情况下，利用 CPU 的 SIMD 资源和向量对齐的跨步技巧 - 所以对用户级编码没有太多期望）

最后但并非最不重要的一点是，有人可能会评论转置的“成本” -numpy对必须转置矩阵没有任何作用，而是改变索引顺序 - 没有别的。如果这可能会产生一些影响，则可以预期，而不是通过查看数据单元的底层存储到物理 RAM 的排序类型FORTRAN或排序语言C类型，但在最大范围内，这使得这方面变得可以忽略不计，并且被具有零回写或其他性能相关影响的处理的缓存内性质很好地掩盖了。1E1 x 1E11E2 x 1E1 x 1E1

结果 ：

鉴于上述所有事实和进一步的观察，真正获得此处定义的计算的更高吞吐量的最便宜和最合理的步骤是在这里移动到线性工作自由度 - [GHz]CPU 芯片越高越好（此处性能线性增长）还拥有尽可能多的 AVX-512 寄存器和尽可能大的 L1i + L1d 缓存（亲和映射避免任何其他 O/S 噪声的策略对于 HPC 级性能目标是显而易见的) 并且依赖已经智能的numpy工具，针对这种用于矩阵处理的 CPU 资源组合进行了微调（如果需要超越float64 IEEE-754表示，另一个故事开始了）。

不要期望用户级代码做得比这更好，numpy-native SIMD-aligned 处理可以并且将会提供。

对于上述给定的规模，内联装配可能会获得优势，但必须以巨大的人力成本来制作和测试这样一个终极方案，但解决方案的概念却是一个神秘的变化。如果市场确实需要这样做，请告诉我。

推荐阅读