避免 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) 我得到的系数约为 40。

一些解释：

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

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

避免 Python 中的嵌套 for 循环

Avoiding nested for loops in Python

python

numpy

matrix

nested-loops