你可以通过将中间存储减少到仅对角元素来改进@Bill的解决方案:
from numpy.core.umath_tests import inner1d
m, n = 1000, 500
a = np.random.rand(m, n)
b = np.random.rand(n, m)
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
np.einsum
会稍微占优势,因为它不需要在加法之前存储所有对角线元素,而是保持一个运行总和。并且它使用SIMD指令,所以你应该比inner1d
获得2或4倍的性能提升。但是在我的系统上,它们的表现也完全相同,即使对于更大的数据也是如此。 - Jaimeumath_tests
是一个稳定的公共 API 吗?这个名称听起来像是私有的,并且似乎比 NumPy 的其他部分文档化程度更低。 - amcnabbinner1d
将被包含在numpy.linalg._umath_linalg
中,不确定它是否会留在numpy.core.umath_tests
中。它可能会移动,但我认为有明确的意图保留它并越来越多地暴露它。 - Jaimeinner1d
的源代码中看到了哪些使用SSE
的地方?如果没有#include "emmintrin.h"
或类似的东西,那么你就无法访问SSE
函数。请参考 einsum.c.src 开头的部分进行比较。 - Jaime