将元素和矩阵乘法与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)])

但是此实现是slow。稍微快一点的方法如下所示：

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进行测试。

btw，请注意：

C = np.dot(A, B)[:, :, 0, :]
没有给出正确的结果。因为我首先通过将结果与此np.dot命令给出的结果进行比较来检查我的方法，这使我受骗了。

Answer 2

我的项目中也有同样的问题。我能想到的最好的方法是，我认为它比使用vstack快一点（也许10％）：

Answer 3

非常灵活，紧凑和快速的解决方案：