Numpy 数组和 for 循环：如何改进它

Question

我是 Python 的新手，我有一个关于 for 循环加速的问题。

让“u”是一个维度为 (N,K) 的 numpy 数组，让“kernel_vect”是一个维度为 (K,) 的 numpy 数组，这两个数组都是 float64 数字。 我想加快下面的代码（例如通过消除 for 循环）

Kernel_appo = np.zeros((N**2,))
    for k in range(K):
        uk = u[:,k]
        Mat_appo = np.outer(uk,uk)
        Kernel_appo = Kernel_appo  + kernel_vect[k] * routines.vec(Mat_appo)

任何想法？ 谢谢！

Answer 1

这是一个更快的实现，没有对 k 的循环：

# Version 2
Kernel_appo = np.zeros((N**2,))
for n1 in range(N):
    for n2 in range(N):
        Kernel_appo[n1*N+n2] = (u[n1,:] * u[n2,:] * kernel_vect).sum()

我们可以通过使用 u 乘积的对称性使其更快：

# Version 3
Kernel_appo = np.zeros((N,N))
for n1 in range(N):
    for n2 in range(n1,N):
        Kernel_appo[n1,n2] = (u[n1,:] * u[n2,:] * kernel_vect).sum()
Kernel_appo = np.triu(Kernel_appo, 1) + np.tril(Kernel_appo.transpose(), 0) # make the matrix symmetric
Kernel_appo = np.ravel(Kernel_appo, order='C')

这是一个删除循环之一的版本：

# Version 4
Kernel_appo = np.zeros((N,N))
for n1 in range(N):
    Kernel_appo[n1,n1:N] = ((u[n1,:] * kernel_vect) * u[n1:N,:]).sum(axis=1)
Kernel_appo = np.triu(Kernel_appo, 1) + np.tril(Kernel_appo.transpose(), 0)
Kernel_appo = np.ravel(Kernel_appo, order='C')

我们仍然有一个关于 N 的循环。但是，由于 N 很小，因此保留它似乎是合理的。 删除它肯定会迫使 numpy 在 memory 中创建巨大的矩阵，这会导致性能下降（如果 N 和 K 非常大，甚至会崩溃）。

请注意，如果 K 更大，版本 4 可能不会那么快（因为临时 numpy 矩阵无法放入 CPU 的缓存中）。

更新：我只是发现在这种情况下可以使用很棒的np.einsum ：

# Version 5
Kernel_appo = np.ravel(np.einsum('ji,ki,i->jk', u, u, kernel_vect, optimize=True), order='C')

做好准备，因为这个更简单的实现也快得多（因为 numpy 能够向量化代码并并行运行）。

以下是我的机器上 N=50 和 K=5000 的性能结果：

Initial code: 58.15 ms
Version 2:    19.94 ms
Version 3:    10.11 ms
Version 4:     5.08 ms
Version 5:     0.57 ms

最终的实现现在比最初的快大约 100 倍！