与 Python For 循环速度作斗争

Question

在我开始之前，我会说我知道以前有人问过这个问题，但我一直在努力实现所建议的方法（例如运行通过 PyPy 实现它）。这是加速代码的最后尝试。

基本上我有一段大约 600 行长的代码。大部分代码需要大约 30 秒才能完成运行，但是一小部分（4 行长）需要 5-15 分钟才能完成运行。原因很简单，它是一个 for 循环中的数学方程式，在 for 循环中，在 for 循环中。所以这个方程式被计算了大约 5000 万次。我承认这需要一段时间，但是当同样的事情在 MATLAB 中是运行时，它通常在一分钟内完成。我相信这是因为 JIT 加速；但我可能是错的。无论哪种方式，这让我觉得必须有一种方法可以加快速度。代码部分如下（所使用的矩阵非常大，所以我想我只是说它们的维度，因为它们中的数字可能会有所不同）。

    for k in range(7500):                   
        for jj in range(2):
            for ii in range(k+1):
                 Y[k][jj,0] += S[ii][jj] * c[k-ii][jj,jj] * U[ii][jj,jj]

其中矩阵（/数组）的大小为：

numpy.shape(Y) = (7500, 2, 2)
numpy.shape(S) = (7500, 2, 1)
numpy.shape(c) = (7500, 2, 2)
numpy.shape(U) = (7500, 2, 2)

有人看到我可以做些什么来加快速度吗？

编辑 1：

根据要求，这里是上面的 MATLAB 版本：

for k=1:7500
    for j=1:2
       for i=1:7500

           Y(j,1,k)=Y(j,1,k)+S(j,1,i)*c(j,j,k+1-i)*U(j,j,i);

       end
    end
end

编辑 2：

应该加了，我用的是3.4.2

此外，遗憾的是我没有代码背后的源数学。我有大约 2/3 的代码，但没有后三分之一。我只有要转换的 MATLAB 代码。（至少现在）

Answer 1

使用np.convolve即可得到结果。

import numpy as np

S = np.random.rand(1000, 2, 1)
c = np.random.rand(1000, 2, 2)
U = np.random.rand(1000, 2, 2)

Y = np.zeros_like(U)
for k in range(1000):
    for jj in range(2):
        for ii in range(k+1):
            Y[k,jj,0] += S[ii,jj,0] * c[k-ii,jj,jj] * U[ii,jj,jj]

Yx = np.zeros_like(Y)
for jj in range(2):
    Yx[:,jj,0] += np.convolve(S[:,jj,0] * U[:,jj,jj], c[:,jj,jj], mode='full')[:Yx.shape[0]]

print(abs(Y - Yx).max())
# -> 3.12638803734e-13

如何找到这个？请注意，事物只是沿 jj 轴相乘，并且 ii 求和实际上是卷积。然后只需在 numpy 函数中正确调整索引即可。

如果您想要更高的速度，将 convolve 替换为 scipy.signal.fftconvolve 可能会加快速度。一些时间：

for loops:         77 s
np.convolve:       33.6 ms
fftconvolve:       1.48 ms

这提供了不错的 ~ 50000 倍加速。

另请注意，您应该始终编写 Y[k,jj,0] 而不是 Y[k][jj,0] —— 因为没有 JIT，后者会创建一个临时数组视图，如果您评估多次表达。将 for 循环表达式中的行重写为

Y[k,jj,0] += S[ii,jj,0] * c[k-ii,jj,jj] * U[ii,jj,jj]

已经将计算速度提高了 4 倍 (!)。

与 Python For 循环速度作斗争

Struggling With Python For-Loop Speed

python

numpy