在NumPy中将元素方式和矩阵乘法与多维数组相结合

Question

我有两个多维NumPy数组,A并且B,A.shape = (K, d, N)和B.shape = (K, N, d).我想在轴0(K)上执行逐元素操作,该操作是在轴1和2(d, N和N, d)上的矩阵乘法.所以结果应该是一个多维数组C用C.shape = (K, d, d),这样C[k] = np.dot(A[k], B[k]).一个天真的实现看起来像这样:

C = np.vstack([np.dot(A[k], B[k])[np.newaxis, :, :] for k in xrange(K)])

但这种实施很慢.稍微快一点的方法看起来像这样:

C = np.dot(A, B)[:, :, 0, :]

它使用np.dot多维数组的默认行为,给我一个形状的数组(K, d, K, d).但是,此方法计算所需的答案K时间(沿轴2的每个条目都相同).渐渐地,它将比第一种方法慢,但开销要小得多.我也知道以下方法:

from numpy.core.umath_tests import matrix_multiply
C = matrix_multiply(A, B)

但我不能保证这个功能可用.因此,我的问题是,NumPy是否提供了一种有效执行此操作的标准方法？一般来说,适用于多维数组的答案是完美的,但仅针对这种情况的答案也是很好的.

编辑:正如@Juh_所指出的,第二种方法是错误的.正确的版本是:

C = np.dot(A, B).diagonal(axis1=0, axis2=2).transpose(2, 0, 1)

但是增加的开销使得它比第一种方法慢,即使对于小矩阵也是如此.对于小型和大型矩阵,最后一种方法是在我的所有时序测试中获胜.我现在正在考虑使用这个,如果没有更好的解决方案,即使这意味着将numpy.core.umath_tests库(用C语言编写)复制到我的项目中.

Answer 1

您的问题的一个可能的解决方案是：

C = np.sum(A[:,:,:,np.newaxis]*B[:,np.newaxis,:,:],axis=2)

然而：

1. 仅当 K 远大于 d 和 N 时，它才比 vstack 方法更快
2. 它们可能是一些内存问题：在上述解决方案中，分配了一个 KxdxNxd 数组（即求和之前的所有可能的乘积对）。实际上，我无法使用大 K、d 和 N 进行测试，因为我的内存不足。

顺便说一句，请注意：

C = np.dot(A, B)[:, :, 0, :]

没有给出正确的结果。它让我受骗了，因为我首先通过将结果与 np.dot 命令给出的结果进行比较来检查我的方法。