python - 如何在 Python 中使 for 循环更快？

Question

我正在用 Python 2.7 编写一个脚本来训练神经网络。作为主脚本的一部分，我需要一个解决二维热传导偏导数计算的程序。以前我用 Fortran 编写了这个程序，然后用 Python 重写了它。 Fortran 需要的时间是 0.1 秒，而 Python 需要 13 秒!这对我来说是绝对不能接受的，因为在那种情况下，计算时间将由解决 PDE 的程序部分决定，而不是由神经网络训练的时期决定。

如何解决这个问题？

似乎我无法向量化矩阵，因为新元素 t[i.j] 是使用值 t[i-1,j] 等计算的。

下面是运行缓慢的部分代码:

while (norm > eps):
    # old value
    t_old = np.copy(t)

    # new value
    for i in xrange(1,n-1):
        for j in xrange(1,m-1):
            d[i] = 0.0
            a[i+1,j] = (0.5*dx/k[i,j] + 0.5*dx/k[i+1,j])
            a[i-1,j] = (0.5*dx/k[i,j] + 0.5*dx/k[i-1,j])
            a[i,j+1] = (0.5*dy/k[i,j] + 0.5*dy/k[i,j+1])
            a[i,j-1] = (0.5*dy/k[i,j] + 0.5*dy/k[i,j-1])
            a[i,j] = a[i+1,j] + a[i-1,j] + a[i,j+1] + a[i,j-1]
            sum = a[i+1,j]*t[i+1,j] + a[i-1,j]*t[i-1,j] + a[i,j+1]*t[i,j+1] + a[i,j-1]*t[i,j-1] + d[i]
            t[i,j] = ( sum + d[i] ) / a[i,j]
            k[i,j] = k_func(t[i,j])
    # matrix 2nd norm
    norm = np.linalg.norm(t-t_old)

Answer 1

纯 Python 优化

这不会带来那么多，但它是最简单的。

• 消除无用代码。在内循环中，d[i] 设置为零。然后在两个地方将它添加到其他东西中。添加 0 不会改变任何内容，因此您可以完全删除 d[i]。
• 只计算一次。 k[i,j]、0.5*dx 和 0.5*dy使用了四次。所以计算一次并将它们分配给局部变量。
• 删除不必要的数组访问。在内部循环中，仅计算和使用a 矩阵的五个元素。因此，用局部变量 a1 替换矩阵元素，直到并包括 a5。

代码现在看起来像这样:

while (norm > eps):
    # old value
    t_old = np.copy(t)

    # new value
    for i in xrange(1,n-1):
        for j in xrange(1,m-1):
            px = 0.5*dx
            py = 0.5*py
            q = k[i,j]
            a1 = (px/q + px/k[i+1,j])
            a2 = (px/q + px/k[i-1,j])
            a3 = (py/q + py/k[i,j+1])
            a4 = (py/q + py/k[i,j-1])
            a5 = a1 + a2 + a3 + a4
            sum = a1*t[i+1,j] + a2*t[i-1,j] + a3*t[i,j+1] + a4*t[i,j-1]
            t[i,j] = sum / a5
            k[i,j] = k_func(t[i,j])
    # matrix 2nd norm
    norm = np.linalg.norm(t-t_old)

由于您的示例未提供完整的工作代码，因此我无法衡量效果。

但是，Python 中的循环效率相对较低。为了在纯 Python 中 获得良好的性能，最好使用列表理解而不是循环。这是因为在理解中，循环是在 C 语言的 Python 运行时中完成的，而不是在 Python 字节码中。但由于我们已经在这里处理 numpy 数组，所以我不会对此展开。

重新编码您的算法以使用 numpy 而不是循环

numpy 背后的基本思想是它为数组操作优化了例程(用 C 或 Fortran 编写)。因此，对于数组操作，您应该使用 numpy 函数而不是循环!

您的循环主要包括用从移动一列或一行的另一个矩阵派生的值填充一个矩阵。为此你可以做这样的事情。

在这个例子中，我将把 k 向下移动一行:

In [1]: import numpy as np

In [2]: k = np.arange(1, 26).reshape([5,5])
Out[2]: 
array([[ 1,  2,  3,  4,  5],
       [ 6,  7,  8,  9, 10],
       [11, 12, 13, 14, 15],
       [16, 17, 18, 19, 20],
       [21, 22, 23, 24, 25]])

In [3]: dx = 0.27
Out[3]: 0.27

In [4]: 0.5*dx/k[1:,:]
Out[4]: 
array([[0.0225    , 0.01928571, 0.016875  , 0.015     , 0.0135    ],
       [0.01227273, 0.01125   , 0.01038462, 0.00964286, 0.009     ],
       [0.0084375 , 0.00794118, 0.0075    , 0.00710526, 0.00675   ],
       [0.00642857, 0.00613636, 0.00586957, 0.005625  , 0.0054    ]])

In [5]: np.insert(0.5*dx/k[1:,:], 0, 0, axis=0)
Out[5]: 
array([[0.        , 0.        , 0.        , 0.        , 0.        ],
       [0.0225    , 0.01928571, 0.016875  , 0.015     , 0.0135    ],
       [0.01227273, 0.01125   , 0.01038462, 0.00964286, 0.009     ],
       [0.0084375 , 0.00794118, 0.0075    , 0.00710526, 0.00675   ],
       [0.00642857, 0.00613636, 0.00586957, 0.005625  , 0.0054    ]])