Python 四维numpy阵列上的矩阵乘法
我需要在两个4D数组(m&n)上执行矩阵乘法,m&n的维数分别为2x2x2x2和2x3x2x2,这将产生一个2x3x2x2x2数组。经过大量的研究(主要是在这个网站上),似乎这可以通过np.einsum或np.tensordot有效地完成,但我无法复制我从Matlab得到的答案(手工验证)。我理解在2D阵列上执行矩阵乘法时这些方法(einsum和tensordot)是如何工作的(解释得很清楚),但我无法为4D阵列获得正确的轴索引。很明显我错过了什么!我的实际问题涉及两个23x23x3x3复数数组,但我的测试数组是:Python 四维numpy阵列上的矩阵乘法,python,arrays,numpy,matrix,Python,Arrays,Numpy,Matrix,我需要在两个4D数组(m&n)上执行矩阵乘法,m&n的维数分别为2x2x2x2和2x3x2x2,这将产生一个2x3x2x2x2数组。经过大量的研究(主要是在这个网站上),似乎这可以通过np.einsum或np.tensordot有效地完成,但我无法复制我从Matlab得到的答案(手工验证)。我理解在2D阵列上执行矩阵乘法时这些方法(einsum和tensordot)是如何工作的(解释得很清楚),但我无法为4D阵列获得正确的轴索引。很明显我错过了什么!我的实际问题涉及两个23x23x3x3复数数组
a = np.array([[1, 7], [4, 3]])
b = np.array([[2, 9], [4, 5]])
c = np.array([[3, 6], [1, 0]])
d = np.array([[2, 8], [1, 2]])
e = np.array([[0, 0], [1, 2]])
f = np.array([[2, 8], [1, 0]])
m = np.array([[a, b], [c, d]]) # (2,2,2,2)
n = np.array([[e, f, a], [b, d, c]]) # (2,3,2,2)
我意识到复杂的数字可能会带来更多的问题,但现在,我只是想了解索引是如何与einsum&Tensordo合作的。我要寻找的答案是这个2x3x2x2阵列:
+----+-----------+-----------+-----------+
| | 0 | 1 | 2 |
+====+===========+===========+===========+
| 0 | [[47 77] | [[22 42] | [[44 40] |
| | [31 67]] | [27 74]] | [33 61]] |
+----+-----------+-----------+-----------+
| 1 | [[42 70] | [[24 56] | [[41 51] |
| | [10 19]] | [ 6 20]] | [ 6 13]] |
+----+-----------+-----------+-----------+
我最近的尝试是使用np.tensordot:
mn = np.tensordot(m,n, axes=([1,3],[0,2]))
这给了我一个2x2x3x2数组,数字正确,但顺序不正确:
+----+-----------+-----------+
| | 0 | 1 |
+====+===========+===========+
| 0 | [[47 77] | [[31 67] |
| | [22 42] | [24 74] |
| | [44 40]] | [33 61]] |
+----+-----------+-----------+
| 1 | [[42 70] | [[10 19] |
| | [24 56] | [ 6 20] |
| | [41 51]] | [ 6 13]] |
+----+-----------+-----------+
我还尝试实施了来自的一些解决方案,但没有成功。非常感谢您的最佳选择,如果您有任何关于如何改进的想法,我们将不胜感激,因为您的缩减维度既不匹配(这将允许广播),也不是“内部”维度(这将与
np.tensordot
)使用np.einsum
np.einsum('ijkl,jmln->imkn', m, n)
array([[[[47, 77],
[31, 67]],
[[22, 42],
[24, 74]],
[[44, 40],
[33, 61]]],
[[[42, 70],
[10, 19]],
[[24, 56],
[ 6, 20]],
[[41, 51],
[ 6, 13]]]])
您可以简单地交换
tensordot
结果上的轴,这样我们仍然可以利用BLAS
与tensordot
-
np.tensordot(m,n, axes=((1,3),(0,2))).swapaxes(1,2)
或者,我们可以在tensordot
调用和转置中交换m
和n
的位置,以重新排列所有轴-
np.tensordot(n,m, axes=((0,2),(1,3))).transpose(2,0,3,1)
通过手动调整和交换轴的形状,我们可以将2D
矩阵乘法与np.dot
相结合,如下所示-
m0,m1,m2,m3 = m.shape
n0,n1,n2,n3 = n.shape
m2D = m.swapaxes(1,2).reshape(-1,m1*m3)
n2D = n.swapaxes(1,2).reshape(n0*n2,-1)
out = m2D.dot(n2D).reshape(m0,m2,n1,n3).swapaxes(1,2)
运行时测试-
将输入数组缩放为10x
形状:
In [85]: m = np.random.rand(20,20,20,20)
In [86]: n = np.random.rand(20,30,20,20)
# @Daniel F's soln with einsum
In [87]: %timeit np.einsum('ijkl,jmln->imkn', m, n)
10 loops, best of 3: 136 ms per loop
In [126]: %timeit np.tensordot(m,n, axes=((1,3),(0,2))).swapaxes(1,2)
100 loops, best of 3: 2.31 ms per loop
In [127]: %timeit np.tensordot(n,m, axes=((0,2),(1,3))).transpose(2,0,3,1)
100 loops, best of 3: 2.37 ms per loop
In [128]: %%timeit
...: m0,m1,m2,m3 = m.shape
...: n0,n1,n2,n3 = n.shape
...: m2D = m.swapaxes(1,2).reshape(-1,m1*m3)
...: n2D = n.swapaxes(1,2).reshape(n0*n2,-1)
...: out = m2D.dot(n2D).reshape(m0,m2,n1,n3).swapaxes(1,2)
100 loops, best of 3: 2.36 ms per loop
为了证明广播也是有效的:
(m[:, :, None, :, :, None] * n[None, :, :, None, :, :]).sum(axis=(1,4))
但发布的其他解决方案可能更快,至少对于大型阵列而言是如此。在APL(一种编程语言)中,矩阵乘法类运算符被推广用于高维对象。乘法/加法发生在左对象的最后一个维度和第二个对象的第一个维度之间。对于2x2x2+.x 2x3x2x2,将删除“内部”(…x2,2x…)维度,结果将是一个2x2x3x2x2对象。我想知道Matlab是否遵循同样的规则。我可以看到某种轴交换可以做到这一点,但我没有走这条路,因为我认为tensordot可以做到这一点,而无需重新排列。为什么einsum和tensordot之间的速度有这么大的差异?@AndrewForbes为了了解
tensordot
如何“展开”剩余的轴,这里有一个相关的帖子-。我认为这种方法最容易应用/理解,但与@Divakar的tensordot解决方案相比,它的速度要慢得多,我感到惊讶。我看到的所有信息都表明einsum将是最有效的方法。@AndrewForbes只是基于BLAS的tensordot
太有效了。是的,不幸的是np。einsum
,虽然对了解爱因斯坦求和符号的人来说非常清楚,但并没有像各种*dot
操作符那样优化。但由于你的玩具问题似乎与实际问题的形式不同,我认为过早的优化不值得。