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

【问题标题】：What is the best way to compute the trace of a matrix product in numpy?在numpy中计算矩阵乘积的最佳方法是什么？
【发布时间】：2013-09-22 03:52:55
【问题描述】：

如果我有 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 数组的矩阵乘积的轨迹？最快或最惯用的方法是什么？

【问题讨论】：

标签： python numpy matrix

【解决方案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 稍慢，但它可能不适用于所有系统，请参阅this question：

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

【讨论】：

在我的机器上，inner1d 和 einsum 方法在速度上基本没有区别。
有人会认为np.einsum 会有一点优势，因为它不必在添加它们之前存储所有对角线元素，它保持一个运行总和。它使用 SIMD 指令，因此您应该比inner1d 提高 2 或 4 倍。但它们在我的系统上的性能也相同，即使对于更大的数据也是如此。
顺便问一下，umath_tests 是稳定的公共 API 吗？这个名字听起来很私密，而且它的文档似乎比 numpy 的其他部分要少。
@amcnabb 这很有趣——你使用的是什么版本的 numpy？检查源代码，inner1d 是一个 C++ 定义，它是为使用SSE 而编写的。见here。可以帮助回答einsum 的问题。
@amcnabb 在 numpy 1.8 中 inner1d 将包含在 numpy.linalg._umath_linalg 中，不确定是否会保留在 numpy.core.umath_tests 中。它可能会四处走动，但我认为有明确的意图保留它并更多地暴露它。

【解决方案2】：

从wikipedia，您可以使用 hadamard 乘积（逐元素乘法）计算迹线：

# 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

【讨论】：

这似乎比 numpy.trace 更快，即使矩阵不是对称的（例如，如果 A 是 1000x100 而B 是 100x1000，反之亦然）。
您可能想提一下，A 和 B 必须是 ndarrays，这样才不会混淆。还应该注意的是，时间很大程度上受您的 numpy 链接到哪种 BLAS 的影响。要提高速度，请考虑使用表达式 np.einsum('ij,ji->',A,B)。
@Ophion 在阅读本文之前已将其包含在我的答案中...我们可能还有另一个案例，说明我的系统上 np.einsum 的神秘缓慢...

【解决方案3】：

请注意，一种轻微的变体是采用vectorized 矩阵的点积。在 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

【讨论】：