通過3-d矩陣的每個切片將2-d矩陣的每列相乘的更有效方式

Question

我有一個8x8x25000陣列W和一個8 x 25000陣列r。 我希望通過r的每列（8x1）將每個8x8切片的W復用多個，並將結果保存在Wres中，這將最終成為8x25000矩陣。

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

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

但這很慢，我希望有更快的方法來實現這一目標。

有任何想法嗎？

Answer 1

只要2個陣列共享相同的1軸長度，Matmul就可以傳播。 來自文檔：

如果任一參數是ND，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

時間等同於@ Imanol的解決方案，所以請選擇兩者。 兩者都比循環快30倍。 在這里，由於陣列的大小， einsum將具有競爭力。 如果陣列大於這些，它可能會勝出，而對於較小的陣列則會輸。 請參閱此討論了解更多信息

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

貸到：@Divakar