Numpy 基于cython的多向量外积并行化

假设有一个形状为（N，m）的numpy数组m，我想计算

res = np.zeros((M, M))
for i in range(N):
   res += np.outer(m[i], m[i])

使用

einsum

可以使该循环更有效，即

res = np.sum(np.einsum('ij,ik->ijk', m , m, axis=0)

但这需要存储一个nxm矩阵，这可能（在我的情况下）非常苛刻

我想用cython构建这个函数，使用并行化

import numpy as np
cimport numpy as np
from cython.parallel import prange

def get_s(double[:,:] m):
    cdef Py_ssize_t i = 0
    cdef int n = m.shape[1]
    
    res = 0.
    for i in prange(n, nogil=True):
       res += np.outer(m[i], m[i])
    return res  

该代码的思想是

---

运行这段代码会产生很多错误，因为我使用的是python对象、不允许的操作，并且我不知道如何正确初始化

res

您不能使用numpy函数（

np.outer

）是nogil上下文。所以你只要用循环把它拼出来。 此外，您的

res

变量似乎是一个数组，因此需要声明一个并初始化它。 最后，您希望循环编译为C，因此使用类型化的MemoryView。使用numpy阵列进行内存管理并获取它们的MemoryView是最简单的。综上所述

%%cython -a

cimport cython

import numpy as np

@cython.boundscheck(False)
@cython.wraparound(False)
def m_outer(double[:, ::1] a):
    n, m = a.shape[0], a.shape[1]
    cdef double[:, ::1] resm = np.zeros((m, m))

    for i in range(a.shape[0]):
        for j in range(a.shape[1]):
            for k in range(a.shape[1]):
                resm[j, k] += a[i, j] * a[i, k]
    return np.asarray(resm)

编写这些东西的一种方法（也许是这种方法）是用python编写（不管速度如何），在一个小示例上验证输出（我使用3x4），然后进行cythonize

cython化时，使用

%cython-a

并检查生成的C代码

现在，这里有两个明显的机会：重新排列循环以提升循环常数和使用prange。两者都留给读者作为练习

最后一个音符。除非是教育性练习，否则请注意，您真正计算的是矩阵积

a.T@a

您的迭代：

In [139]: res = np.zeros((4,4))                                                                      
In [140]: for i in range(3): res += np.outer(m[i],m[i])                                              
In [141]: res                                                                                        
Out[141]: 
array([[ 80.,  92., 104., 116.],
       [ 92., 107., 122., 137.],
       [104., 122., 140., 158.],
       [116., 137., 158., 179.]])

我们可以对广播执行相同的

outer

：

In [142]: np.sum(m[:,:,None]*m[:,None,:], axis=0)                                                    
Out[142]: 
array([[ 80,  92, 104, 116],
       [ 92, 107, 122, 137],
       [104, 122, 140, 158],
       [116, 137, 158, 179]])

（是的，这会生成一个临时（N，M，M）数组）

建议的单步einsum：

In [143]: np.einsum('ij,ik->jk',m,m)                                                                 
Out[143]: 
array([[ 80,  92, 104, 116],
       [ 92, 107, 122, 137],
       [104, 122, 140, 158],
       [116, 137, 158, 179]])

这只是一个简单的点积（带有适当的转置）：

---

由于

numpy

dot使用快速BLAS代码，我怀疑您是否可以使用

cython

对其进行改进。这个问题可能需要一些编辑：没有一个代码片段是完整的或有效的，句子缺少部分。如果您有多个问题，请询问多个问题（在联机搜索并记录您的搜索之后）。您的einsum使用没有多大意义<代码>np.einsum（'ij，ik->jk'，A，A，optimize='optimal'）就足够了。这实际上只是一个点积，如果打开优化，einsum会检测到它。我想这是最好的方法。我一个人没有看到它，我觉得很傻！谢谢：）。

In [144]: m.T.dot(m)                                                                                 
Out[144]: 
array([[ 80,  92, 104, 116],
       [ 92, 107, 122, 137],
       [104, 122, 140, 158],
       [116, 137, 158, 179]])
In [145]: m.T@m                                                                                      
Out[145]: 
array([[ 80,  92, 104, 116],
       [ 92, 107, 122, 137],
       [104, 122, 140, 158],
       [116, 137, 158, 179]])