避免在Python中嵌套for循环

Posted

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了避免在Python中嵌套for循环相关的知识,希望对你有一定的参考价值。

我有一个矩阵A维度1024 * 307200和另一个矩阵B维度1024 * 50。我在嵌套的for循环中对这两个矩阵执行L2_norm,以使我的最终矩阵C307200 * 50。 您可以在下面找到以下代码:

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

如您所见,我的变量的维度很大,导致非常高的延迟。我想避免这个嵌套循环,因为对于il的每个值,我使用相同的函数。

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

答案

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

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)来节省记忆。

以上是关于避免在Python中嵌套for循环的主要内容,如果未能解决你的问题,请参考以下文章

python中为啥我的for循环里嵌套的if只能循环一次?

一文了解Python中的循环(for while break continue 嵌套循环...)

python之控制代码if,for,while

嵌套for循环以在Python中递归

在 Python 的嵌套 for 循环中使用 Continue 语句

嵌套for循环慢python用于计算特殊标准偏差