3d 矩阵乘法到多个 2d 矩阵乘法

Question

我有两个 3d 矩阵 X 和 Y，它们的形状都是 (5, 1825, 77)。我想做五次二维矩阵乘法，即 X[i, :, :]@Y[i, :, :].T 而不使用 for 循环。有没有办法在 numpy 中做到这一点？

Answer 1

对于那些不惜一切代价避免 for 循环的人（比如我）来说，这很有趣：

shape = 5, 1825, 77
X = np.random.random(shape)
Y = np.random.random(shape)

p_for = np.empty((shape[0], shape[1], shape[1]))
for i in range(shape[0]):
    p_for[i] = X[i] @ Y[i].T

p_matmul = X @ np.moveaxis(Y, -1, 1)
assert np.allclose(p_for, p_matmul)

p_einsum = np.einsum("ijk,ilk->ijl", X, Y)
assert np.allclose(p_for, p_einsum)

树方法产生相同的结果，但是，正如@JérômeRichard 指出的那样：

%%timeit
prod = np.empty((shape[0], shape[1], shape[1]))
for i in range(5):
    prod[i] = X[i, :, :] @ Y[i, :, :].T
50.4 ms ± 7.18 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit X @ np.moveaxis(Y, -1, 1)
115 ms ± 1.6 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit np.einsum("ijk,ilk->ijl", X, Y)
544 ms ± 3.22 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

3d 矩阵乘法到多个 2d 矩阵乘法

3d matrix multiplication into multiple 2d matrix multiplications

numpy

matrix-multiplication