NumPy:将每对二维矩阵乘以两个三维矩阵的有效方法

Question

我在NumPy中有两个3d矩阵，它们的高度、宽度与矩阵乘法匹配。(16X3和3X16(

x = np.random.random((125,16,3)).astype(np.float32)
y = np.random.random((14,3,16)).astype(np.float32)

我想将x中的每个2d矩阵与y中的每个2d矩阵相乘，这样得到的结果是形状为res[x.shape[0]，y.shape[0]…]的4d矩阵。目前，我正在使用此代码。

xd,xh,xw = x.shape
yd,yh,yw = y.shape
res = np.zeros((xd,yd,xh,yw))
for i in range(xd):
for j in range(yd):
res[i,j,...] = x[i,...]@y[j,...]

有没有其他方法可以在没有循环的情况下实现这一点？一些更快的NumPy方式？

Answer 1

看看numpy广播

import numpy as np
l, k, n, m = 11, 13, 17, 19
x = np.random.random((l, n, m))
y = np.random.random((k, m, n))

# (l, k, n, n) =  (l,          1, n, m) @ (          1, k, m, n)
z              = x[:, np.newaxis, :, :] @ y[np.newaxis, :, :, :]