由于大量的Numpy点调用而使开销最小化

Question

我的问题如下，我有一个迭代算法，因此在每次迭代时，对于i = 1 ... k，它都需要执行几个矩阵-矩阵乘法dot（A_i，B_i） 。由于这些乘法是通过Numpy的点执行的，因此我知道它们正在调用BLAS-3实现，这是非常快的。问题在于调用的数量巨大，而事实证明这是我程序中的瓶颈。我想通过生产较少的产品但使用更大的矩阵来最大程度地减少所有这些调用的开销。

为简单起见，请考虑所有矩阵均为n x n（通常n不大，范围在1到1000之间）。解决我的问题的一种方法是考虑块对角矩阵diag（A_i）并执行以下乘积。

这只是对函数点的一次调用，但是现在程序浪费很多时间执行与零的乘法运算。这个想法似乎没有用，但它给出了结果[A_1 B_1，...，A_k B_k]，即所有产品堆叠在一个大矩阵中。

我的问题是这样，有没有一种方法可以通过单个函数调用来计算[A_1 B_1，...，A_k B_k]？甚至更重要的是，与环回Numpy点相比，我如何能更快地计算这些乘积？

Answer 1

您可以将数组堆叠为形状（k，n，n），并调用numpy.matmul或使用@运算符。

例如，

In [18]: A0 = np.array([[1, 2], [3, 4]])                                                                 

In [19]: A1 = np.array([[1, 2], [-3, 5]])                                                                

In [20]: A2 = np.array([[4, 0], [1, 1]])                                                                 

In [21]: B0 = np.array([[1, 4], [-3, 4]])                                                                

In [22]: B1 = np.array([[2, 1], [1, 1]])                                                                 

In [23]: B2 = np.array([[-2, 9], [0, 1]])                                                                

In [24]: np.matmul([A0, A1, A2], [B0, B1, B2])                                                           
Out[24]: 
array([[[-5, 12],
        [-9, 28]],

       [[ 4,  3],
        [-1,  2]],

       [[-8, 36],
        [-2, 10]]])

或者，使用@：

In [32]: A = np.array([A0, A1, A2])                                                                      

In [33]: A                                                                                               
Out[33]: 
array([[[ 1,  2],
        [ 3,  4]],

       [[ 1,  2],
        [-3,  5]],

       [[ 4,  0],
        [ 1,  1]]])

In [34]: B = np.array([B0, B1, B2])                                                                      

In [35]: A @ B                                                                                           
Out[35]: 
array([[[-5, 12],
        [-9, 28]],

       [[ 4,  3],
        [-1,  2]],

       [[-8, 36],
        [-2, 10]]])

Answer 2

我似乎您的所有B数组都是一维的，因此@WarrenWeckesser的方法可能会遇到一些问题（因为matmul和@会贪婪地假定两个数组至少为2d）。您可能需要改用np.einsum，它为如何处理广播和尺寸匹配提供了更大的灵活性。

np.einsum

因为点积的

A = [A0, A1, A2, . . . ]   #k (n, n) arrays
B = [b0, b1, b2, . . . ]   #k (n,) arrays

out = np.einsum('kij, kj -> ki', np.array(A), np.array(B))

为Einstein Notation，所以在保持'ij, j -> i'尺寸（这是k和A的“列表”的尺寸）的同时，也可以这样做]

由于这一切都是在C代码中进行的，没有回调python，因此ti应该比B循环快得多