performance - 高效计算 1xM Times MxM Times Mx1, N Times

Question

我需要计算 N 行的点积；假设对于每一行，我需要计算 1xM 乘以 MxM 乘以 Mx1。如果我只查看一行，则计算结果为 1x1。然而，我有 N 行，所以我想我应该堆叠这些行，并将其输入点积；然而我得到了 NxN 矩阵作为结果。我需要的结果位于对角线中，但是是否有一种更快的方法来进行此计算，并且不浪费空间？理想情况下，我希望最终得到一个 Nx1 向量，而不是 NxN 矩阵。

示例

单行

r = np.array([[1,2]]).T
R = np.array([[2,2],[2,2]])

给出

[[18]]

多行

rs = np.array([[1,2],[4,4]]).T
R = np.array([[2,2],[2,2]])
print np.dot(np.dot(rs.T,R), rs)

[[ 18  48]
 [ 48 128]]

Answer 1

使用np.einsum以及内部点积，就像这样 -

np.einsum('ij,ji->i',np.dot(rs.T,R),rs)

示例运行 -

In [215]: rs = np.random.rand(3,4)
     ...: R = np.random.rand(3,3)
     ...: out = np.dot(np.dot(rs.T,R), rs)
     ...: 

In [216]: np.diag(out)  # Diagonal elems is the expected o/p
Out[216]: array([ 1.11476081,  1.05112902,  0.32136029,  0.31318894])

In [217]: np.einsum('ij,ji->i',np.dot(rs.T,R),rs)
Out[217]: array([ 1.11476081,  1.05112902,  0.32136029,  0.31318894])

运行时测试 -

In [233]: rs = np.random.rand(300,400)

In [234]: R = np.random.rand(300,300)

In [235]: %timeit np.diag(np.dot(np.dot(rs.T,R), rs))# Original soln
10 loops, best of 3: 84 ms per loop

In [236]: %timeit np.einsum('ij,kj,ki->j', rs, rs, R)# @DSM's soln
10 loops, best of 3: 65.2 ms per loop

In [237]: %timeit np.einsum('ij,ji->i',np.dot(rs.T,R),rs)
10 loops, best of 3: 37.2 ms per loop