高效计算 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