将 2-d 矩阵的每一列乘以 3-d 矩阵的每个切片的更有效方法

Question

我有一个 8x8x25000 阵列 W 和一个 8 x 25000 阵列 r。我想将 W 的每个 8x8 切片乘以 r 的每一列 (8x1) 并将结果保存在 Wres 中，这最终将成为一个 8x25000 矩阵。

我正在使用 for 循环来完成此操作：

for i in range(0,25000):
    Wres[:,i] = np.matmul(W[:,:,i],res[:,i])

但这很慢，我希望有更快的方法来完成。

有什么想法吗？

Answer 1

只要 2 个数组共享相同的 1 个轴长度，Matmul 就可以传播。来自文档：

如果任一参数为 N-D，N > 2，则将其视为驻留在最后两个索引中的矩阵堆栈并相应地广播。

因此，您必须在matmul 之前执行2 次操作：

import numpy as np
a = np.random.rand(8,8,100)
b = np.random.rand(8, 100)

1. 转置a 和b 使第一个轴是100 个切片
2. 为b 添加一个额外的维度，以便b.shape = (100, 8, 1)

然后：

 at = a.transpose(2, 0, 1) # swap to shape 100, 8, 8
 bt = b.T[..., None] # swap to shape 100, 8, 1
 c = np.matmul(at, bt)

c 现在是100, 8, 1，重新整形回8, 100：

 c = np.squeeze(c).swapaxes(0, 1)

或

 c = np.squeeze(c).T

---

最后，只为方便起见的单行：

c = np.squeeze(np.matmul(a.transpose(2, 0, 1), b.T[..., None])).T

Answer 2

使用 np.matmul 的替代方法是 np.einsum，它可以在 1 行更短且可以说更可口的代码中完成，无需方法链接。

示例数组：

np.random.seed(123)
w = np.random.rand(8,8,25000)
r = np.random.rand(8,25000)
wres = np.einsum('ijk,jk->ik',w,r)

# a quick check on result equivalency to your loop
print(np.allclose(np.matmul(w[:, :, 1], r[:, 1]), wres[:, 1]))
True

Timing 等同于@Imanol 的解决方案，因此请选择两者。两者都比循环快 30 倍。在这里，einsum 将因为数组的大小而具有竞争力。对于比这些更大的阵列，它可能会胜出，而在较小的阵列上会输。有关更多信息，请参阅this 讨论。

def solution1():
    return np.einsum('ijk,jk->ik',w,r)

def solution2():
    return np.squeeze(np.matmul(w.transpose(2, 0, 1), r.T[..., None])).T

def solution3():
    Wres = np.empty((8, 25000))
    for i in range(0,25000):
        Wres[:,i] = np.matmul(w[:,:,i],r[:,i])
    return Wres

%timeit solution1()
100 loops, best of 3: 2.51 ms per loop

%timeit solution2()
100 loops, best of 3: 2.52 ms per loop

%timeit solution3()
10 loops, best of 3: 64.2 ms per loop

Credit 至：@Divakar