避免在Python中嵌套for循環

Question

我有一個尺寸為1024 * 307200的矩陣A和另一個尺寸為1024 * 50矩陣B 我在嵌套的for循環中對這兩個矩陣執行L2_norm ，以使最終矩陣C為307200 * 50 。
您可以在下面找到代碼：

    for i in range(307200):
        for l in range(50):
            C[i,l] =  numpy.linalg.norm(A[:,i] - B[:,l]))      

如您所見，我的變量的維數很大，這導致非常高的延遲。 我想避免這種嵌套循環，因為對於i和l每個值，我都使用相同的函數。

有什么方法可以優化上述循環嗎？

Answer 1

也許您可以用這些矩陣運算來替換內部循環和函數？

for i in range(307200):
    temp = A[:,i,np.newaxis] - B[:]
    C[i,:] = np.linalg.norm(temp, axis=0)

對於較小的陣列，我的運行時間縮短了約20倍。 也許您會收獲更多。 無論如何，請確保您收到的效果不錯（使用較小的電視機）。

Answer 2

更新：通過OP的更新和說明，事情變得簡單得多：

>>> def f_pp(A, B):
...     return np.sqrt(np.add.outer(np.einsum('ij,ij->j', A, A), np.einsum('il,il->l', B, B)) - 2*np.einsum('ij,il->jl', A, B))
... 

結束更新

您可以使用np.einsum和實數運算來大幅提高速度：

>>> def f_pp(A, B):
...     Ar = A.view(float).reshape(*A.shape, 2)
...     Br = B.view(float).reshape(*B.shape, 2)
...     return np.sqrt(np.add.outer(np.einsum('ijk,ijk->j', Ar, Ar), np.einsum('ilk,ilk->l', Br, Br)) - 2*np.einsum('ijk,ilk->jl', Ar, Br))
... 

對於形狀(1024, 3072)和(1024, 50) (1024, 3072) ，我得到的系數約為40 。

一些解釋：

實數運算：除非numpy進行一些令人難以置信的智能優化，否則我希望像x*x.conj()這樣的復雜乘積使用4個實數乘法。 知道結果是真實的，我們保存其中兩個。

將|AB|^2為|A|^2 + |B|^2 - 2|A*B| 。 這樣可以避免直接廣播方法將使用的巨大中間AB形狀(1024, 3072, 50) (1024, 307200, 50) （在完整示例中為(1024, 307200, 50)節省了內存。