numpy中计算矩阵乘积的迹(trace)的最佳方法是什么？

Question

如果我有numpy数组A和B，那么我可以使用以下代码计算它们矩阵乘积的迹：

```python import numpy as np np.trace(np.dot(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数组的矩阵乘积的迹？最快或最惯用的方法是什么？

- amcnabb

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 运行比使用 einsum 稍慢，但这可能并不适用于所有系统，请参见此问题：

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

- Jaime

---

人们可能会认为np.einsum会稍微占优势，因为它不需要在加法之前存储所有对角线元素，而是保持一个运行总和。并且它使用SIMD指令，所以你应该比inner1d获得2或4倍的性能提升。但是在我的系统上，它们的表现也完全相同，即使对于更大的数据也是如此。 - Jaime

顺便问一句，umath_tests 是一个稳定的公共 API 吗？这个名称听起来像是私有的，并且似乎比 NumPy 的其他部分文档化程度更低。 - amcnabb

@amcnabb 在numpy 1.8中，inner1d将被包含在numpy.linalg._umath_linalg中，不确定它是否会留在numpy.core.umath_tests中。它可能会移动，但我认为有明确的意图保留它并越来越多地暴露它。 - Jaime

@Ophion 你在 inner1d 的源代码中看到了哪些使用 SSE 的地方？如果没有 #include "emmintrin.h" 或类似的东西，那么你就无法访问 SSE 函数。请参考 einsum.c.src 开头的部分进行比较。 - Jaime

是的，这是可能的，并且可以解释我的系统上的差异，因为它已经使用MSVC编译，我不认为会进行这样的优化。 - Jaime

Answer 2

从 维基百科 上可以使用哈达玛积（逐元素相乘）计算迹：

# 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

- wflynny

---

即使矩阵不对称（例如，如果 A 是 1000x100 而 B 是 100x1000 或者反之），这似乎比 numpy.trace 更快。 - amcnabb

你可能需要提到A和B必须是“ndarrays”以避免混淆。此外，需要注意的是计时受到numpy链接的BLAS类型的影响。为了进一步提高速度，请考虑使用表达式“np.einsum（'ij，ji->'，A，B）”。 - Daniel

@Ophion 在阅读这篇文章之前，我已经在我的答案中包含了这个...我们可能又遇到了np.einsum在我的系统上神秘缓慢的情况... - Jaime

Answer 3

请注意，有一种微小的变体是对向量化矩阵进行点积运算。在Python中，可以使用`.flatten('F')`来进行向量化操作。在我的电脑上，这种方法比取Hadamard乘积的和稍慢一些，所以相对于wflynny的解决方案来说，它并不是一个更好的选择。但我认为它很有趣，因为在某些情况下可能更直观。例如，就个人而言，我发现对于矩阵正态分布，向量化的解决方案更容易理解。

在我的系统上进行速度比较：

import numpy as np
import time

N = 1000

np.random.seed(123)
A = np.random.randn(N, N)
B = np.random.randn(N, N)

tart = time.time()
for i in range(10):
    C = np.trace(A.dot(B))
print(time.time() - start, C)

start = time.time()
for i in range(10):
    C = A.flatten('F').dot(B.T.flatten('F'))
print(time.time() - start, C)

start = time.time()
for i in range(10):
    C = (A.T * B).sum()
print(time.time() - start, C)

start = time.time()
for i in range(10):
    C = (A * B.T).sum()
print(time.time() - start, C)

结果：

6.246593236923218 -629.370798672
0.06539678573608398 -629.370798672
0.057890892028808594 -629.370798672
0.05709719657897949 -629.370798672

- Hugh Perkins