避免 Python 中的嵌套 for 循环

Question

我有一个矩阵A，维度为1024 * 307200，另一个矩阵B，维度为1024 * 50。我在嵌套的 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)，我得到的系数约为40。

一些解释：

真正的算术：除非 numpy 进行一些令人难以置信的智能优化，否则我希望像 x*x.conj() 这样的复杂产品使用 4 个实数乘法。知道结果是真实的，我们保存了其中的两个。

将|A-B|^2 写为|A|^2 + |B|^2 - 2|A*B|。这通过避免直接广播方法将使用的形状为(1024, 3072, 50)（完整示例中为(1024, 307200, 50)）的巨大中间A-B 来节省内存。