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)

然而：

只有当K比d和N
它们可能是一些内存问题：在上述解决方案中，分配了KxdxNxd数组（即，在求和之前所有可能的产品对）。实际上我无法用大K，D和N进行测试，因为我内存不足。

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

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

Answer 2

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

K, d, N = A.shape
C = np.empty((K, d, d))
for k in xrange(K):
    C[k] = np.dot(A[k], B[k])

我很想看到一个更好的解决方案，我不太清楚如何使用tensordot来做到这一点。

Answer 3

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

C = np.einsum('Kab,Kbc->Kac', A, B, optimize=True)

确认：

import numpy as np
K = 10
d = 5
N = 3
A = np.random.rand(K,d,N)
B = np.random.rand(K,N,d)
C_old = np.dot(A, B).diagonal(axis1=0, axis2=2).transpose(2, 0, 1)
C_new = np.einsum('Kab,Kbc->Kac', A, B)
print(np.max(C_old-C_new))  # should be 0 or a very small number

对于大型多维数组，可选参数optimize=True可以节省大量时间。您可以在这里了解 einsum ：

https://ajcr.net/Basic-guide-to-einsum/

https://rockt.github.io/2018/04/30/einsum

https://numpy.org/doc/stable/reference/generated/numpy.einsum.html

报价：

爱因斯坦求和约定可用于计算许多多维的线性代数数组运算。 einsum 提供了一种简洁的表示方式。这些操作的详尽列表如下：

数组的跟踪 numpy.trace 。

返回对角线 numpy.diag 。

数组轴总和， numpy.sum 。

转置和置换， numpy.transpose 。

矩阵乘法和点积 numpy.matmul numpy.dot 。

向量内部和外部乘积， numpy.inner 。 numpy.outer 。

广播，逐元素和标量乘法， numpy.multiply 。

张量收缩， numpy.tensordot 。

以有效的计算顺序 numpy.einsum_path 约束了数组操作。

Answer 4

你可以做

np.matmul(A, B)

看看https://numpy.org/doc/stable/reference/generated/numpy.matmul.html。

对于足够大的K，应该比einsum快。

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

4 个答案: