在numpy中计算矩阵乘积轨迹的最佳方法是什么？

Question

如果我有numpy的阵列A和B,那么我可以计算出他们的矩阵产品的跟踪:

tr = numpy.linalg.trace(A.dot(B))

然而,A.dot(B)当在迹线中仅使用对角线元素时,矩阵乘法不必要地计算矩阵乘积中的所有非对角线条目.相反,我可以这样做:

tr = 0.0
for i in range(n):
    tr += A[i, :].dot(B[:, i])

但是这会在Python代码中执行循环,并不像那样明显numpy.linalg.trace.

有没有更好的方法来计算numpy数组矩阵乘积的轨迹？什么是最快或最惯用的方式？

Answer 1

您可以通过将中间存储减少到对角元素来改进@ Bill的解决方案:

from numpy.core.umath_tests import inner1d

m, n = 1000, 500

a = np.random.rand(m, n)
b = np.random.rand(n, m)

# They all should give the same result
print np.trace(a.dot(b))
print np.sum(a*b.T)
print np.sum(inner1d(a, b.T))

%timeit np.trace(a.dot(b))
10 loops, best of 3: 34.7 ms per loop

%timeit np.sum(a*b.T)
100 loops, best of 3: 4.85 ms per loop

%timeit np.sum(inner1d(a, b.T))
1000 loops, best of 3: 1.83 ms per loop

另一种选择是使用np.einsum并且根本没有明确的中间存储:

# Will print the same as the others:
print np.einsum('ij,ji->', a, b)

在我的系统上,它运行速度比使用稍慢inner1d,但它可能不适用于所有系统,请参阅此问题:

%timeit np.einsum('ij,ji->', a, b)
100 loops, best of 3: 1.91 ms per loop

Answer 2

从维基百科您可以使用hammard产品计算跟踪(逐元素乘法):

# Tr(A.B)
tr = (A*B.T).sum()

我认为这比计算更少numpy.trace(A.dot(B)).

编辑:

跑一些计时器.这种方式比使用方式快得多numpy.trace.

In [37]: timeit("np.trace(A.dot(B))", setup="""import numpy as np;
A, B = np.random.rand(1000,1000), np.random.rand(1000,1000)""", number=100)
Out[38]: 8.6434469223022461

In [39]: timeit("(A*B.T).sum()", setup="""import numpy as np;
A, B = np.random.rand(1000,1000), np.random.rand(1000,1000)""", number=100)
Out[40]: 0.5516049861907959