避免在Python中嵌套for循环

Posted 2021-03-28

我有一个矩阵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的每个值，我使用相同的函数。

有没有办法优化上面的循环？

答案

也许你可以用这些矩阵运算替换内循环和函数？

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

对于较小的阵列，我的运行时间减少了约20倍。也许你获得更多。在任何情况下，请确保您获得了良好的结果（在较小的集合上）。

另一答案

更新：随着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|。这通过避免直接广播方法将使用的形状A-B（在完整示例中为(1024, 3072, 50)）的巨大中间(1024, 307200, 50)来节省记忆。