在 Python 中使用嵌套循环重写 MATLAB 代码并快速执行

Oga*_*shi 0 python matlab numpy cython nested-loops

这是一个嵌套循环,其中内部索引取决于外部索引,具有以下参数:

f = rand(1,70299)
nech=24*30*24
N=length(f);
xh=(1:nech)/24;
Run Code Online (Sandbox Code Playgroud)

在 MATLAB 中:

sf2(1:nech)=0.;
sf2vel(1:nech)=0.;
count(1:nech)=0.;

for i=1:nech
    for j=1:N-i-1
        sf2(i)=sf2(i)+(f(j+i)-f(j))^2;
        count(i)=count(i)+1; 
    end
    sf2(i)=sf2(i)/count(i);
end
Run Code Online (Sandbox Code Playgroud)

在Python中:

def structFunPython(f,N,nech):
    sf2 = np.zeros(N)
    count = np.zeros(N)
    for i in range(nech):
        indN = np.arange(1,N-i-1)
        for j in indN:
            sf2[i] += np.power((f[i+j]-f[j]),2)
            count[i] += 1
        sf2[i] = sf2[i]/count[i]
    return sf2
Run Code Online (Sandbox Code Playgroud)

与赛通:

import cython
cimport numpy as np
import numpy as np
def structFun(np.ndarray f,N,nech):
    cdef np.ndarray sf2 = np.zeros(N), count = np.zeros(N),
    for i in range(nech):
        indN = np.arange(1,N-i-1)
        for j in indN:
            sf2[i] += np.power((f[i+j]-f[j]),2)
            count[i] += 1
        sf2[i] = sf2[i]/count[i]
    return sf2
Run Code Online (Sandbox Code Playgroud)

执行次数:

Matlab: 7.8377 sec
Python: 3651.35 sec
Cython: 3336.21 sec
Run Code Online (Sandbox Code Playgroud)

我很难相信 Python 和 Cython (尤其是 Cython)对于相同的计算来说那么慢,所以我想我一定在我的 Python/Cython 循环中犯了错误,但我看不出在哪里。

nor*_*ok2 6

我将 MATLAB 代码重写为等效的 Python 代码的方式可能是(注意 MATLAB 中从 1 开始的索引和 Python 中从 0 开始的索引...因为我不知道在没有上下文的情况下应该如何调整它,所以我采用了最简单的方法):

\n
import numpy as np\n\n\ndef func(f_arr, nech, n):\n    sf2 = np.zeros(nech)\n    count = np.zeros(nech)\n    for i in range(nech):\n        for j in range(n - i):\n            sf2[i] += (f_arr[i + j] - f_arr[i]) ** 2\n            count[i] += 1\n        sf2[i] /= count[i]\n    return sf2\n
Run Code Online (Sandbox Code Playgroud)\n

请注意,count[i] += 1是无用的,因为 的最终值count[i]是预先知道的,实际上整个count都是无用的,例如:

\n
import numpy as np\n\n\ndef func2(f_arr, nech, n):\n    sf2 = np.zeros(nech)\n    for i in range(nech):\n        for j in range(n - i):\n            sf2[i] += (f_arr[i + j] - f_arr[i]) ** 2\n        sf2[i] /= (n - i)\n    return sf2\n
Run Code Online (Sandbox Code Playgroud)\n
\n

加速

\n

这是 Numba 加速的手动案例。这就像添加/使用 Numba@njit装饰器一样简单:

\n
import numba as nb\n\n\nfunc_nb = nb.njit(func)\nfunc2_nb = nb.njit(func2)\n
Run Code Online (Sandbox Code Playgroud)\n

现在,funcfunc_nbfunc2func2_nb执行相同的计算:

\n
nech = n = 6\nf_arr = np.arange(n)\nprint(func(f_arr, nech, n))\n# [9.16666667 6.         3.5        1.66666667 0.5        0.        ]\nprint(func_nb(f_arr, nech, n))\n# [9.16666667 6.         3.5        1.66666667 0.5        0.        ]\nprint(func2(f_arr, nech, n))\n# [9.16666667 6.         3.5        1.66666667 0.5        0.        ]\nprint(func2_nb(f_arr, nech, n))\n# [9.16666667 6.         3.5        1.66666667 0.5        0.        ]\n
Run Code Online (Sandbox Code Playgroud)\n

如果您确实需要坚持使用 Cython,这里有一个基于以下的实现func2

\n
%%cython -c-O3 -c-march=native -a\n#cython: language_level=3, boundscheck=False, wraparound=False, initializedcheck=False, cdivision=True, infer_types=True\n\n\nimport numpy as np\nimport cython as cy\n\n\ncpdef func2_cy(f_arr, nech, n):\n    sf2 = np.zeros(nech)\n    _func2_cy(f_arr.astype(np.float_), sf2, nech, n)\n    return sf2\n\n\ncdef _func2_cy(double[:] f_arr, double[:] sf2, cy.int nech, cy.int n):\n    for i in range(nech):\n        for j in range(1, n - i):\n            sf2[i] = sf2[i] + (f_arr[i + j] - f_arr[i]) ** 2\n        sf2[i] = sf2[i] / (n - i)\n
Run Code Online (Sandbox Code Playgroud)\n

与 Numba 相比,编写起来要复杂得多,但可以实现相似的性能。\n诀窍是拥有一个_func2_cy几乎不与 Python 交互的函数(阅读:它以 C 速度运行)。

\n

结果再次与以下相同func2

\n
%%cython -c-O3 -c-march=native -a\n#cython: language_level=3, boundscheck=False, wraparound=False, initializedcheck=False, cdivision=True, infer_types=True\n\n\nimport numpy as np\nimport cython as cy\n\n\ncpdef func2_cy(f_arr, nech, n):\n    sf2 = np.zeros(nech)\n    _func2_cy(f_arr.astype(np.float_), sf2, nech, n)\n    return sf2\n\n\ncdef _func2_cy(double[:] f_arr, double[:] sf2, cy.int nech, cy.int n):\n    for i in range(nech):\n        for j in range(1, n - i):\n            sf2[i] = sf2[i] + (f_arr[i + j] - f_arr[i]) ** 2\n        sf2[i] = sf2[i] / (n - i)\n
Run Code Online (Sandbox Code Playgroud)\n

时间安排

\n

通过一些小玩具基准测试,我们感受到了加速,包括像您一样编写的类似函数,以及Andras Deak 的非常好的答案中提出的矢量化解决方案(但修复了索引以匹配上述内容):

\n
nech = n = 6\nf_arr = np.arange(n)\nprint(func2_cy(f_arr, nech, n))\n# [9.16666667 6.         3.5        1.66666667 0.5        0.        ]\n
Run Code Online (Sandbox Code Playgroud)\n
def func_OP(f, nech, n):\n    sf2 = np.zeros(n)\n    count = np.zeros(n)\n    for i in range(nech):\n        indN = np.arange(n - i)  # <-- indexing fixed\n        for j in indN:\n            sf2[i] += np.power((f[i+j]-f[i]),2)\n            count[i] += 1\n        sf2[i] = sf2[i] / count[i]\n    return sf2\n\n\nfunc_OP_nb = nb.njit(func_OP)\n
Run Code Online (Sandbox Code Playgroud)\n
def func_pdisty(f, nech, n):\n    res = np.zeros(nech)\n    dists = scipy.spatial.distance.pdist(f[:, None], metric=\'sqeuclidean\')\n    counts = np.arange(n - 1, n - nech - 1, -1)\n    inds = np.repeat(np.arange(res.size), counts)\n    np.add.at(res, inds, dists[:inds.size])\n    res /= (counts + 1)\n    return res\n
Run Code Online (Sandbox Code Playgroud)\n
nech = n = 6\nf_arr = np.arange(n)\nprint(func_OP(f_arr, nech, n))\n# [9.16666667 6.         3.5        1.66666667 0.5        0.        ]\nprint(func_pdisty(f_arr, nech, n))\n# [9.16666667 6.         3.5        1.66666667 0.5        0.        ]\n
Run Code Online (Sandbox Code Playgroud)\n
nech = n = 1000\nf_arr = np.arange(n)\n%timeit func_OP(f_arr, nech, n)\n# 1 loop, best of 5: 1.5 s per loop\n%timeit func(f_arr, nech, n)\n# 1 loop, best of 5: 567 ms per loop\n%timeit func2(f_arr, nech, n)\n# 1 loop, best of 5: 352 ms per loop\n%timeit func_OP_nb(f_arr, nech, n)\n# 1000 loops, best of 5: 1.87 ms per loop\n%timeit func_nb(f_arr, nech, n)\n# 1000 loops, best of 5: 1.7 ms per loop\n%timeit func2_nb(f_arr, nech, n)\n# 1000 loops, best of 5: 768 \xc2\xb5s per loop\n%timeit func_pdisty(f_arr, nech, n)\n# 10 loops, best of 5: 44.5 ms per loop\n%timeit func2_cy(f_arr, nech, n)\n# 1000 loops, best of 5: 1 ms per loop\n
Run Code Online (Sandbox Code Playgroud)\n
nech = n = 2000\nf_arr = np.arange(n)\n%timeit func_OP(f_arr, nech, n)\n# 1 loop, best of 5: 6.01 s per loop\n%timeit func(f_arr, nech, n)\n# 1 loop, best of 5: 2.3 s per loop\n%timeit func2(f_arr, nech, n)\n# 1 loop, best of 5: 1.42 s per loop\n%timeit func_OP_nb(f_arr, nech, n)\n# 100 loops, best of 5: 7.31 ms per loop\n%timeit func_nb(f_arr, nech, n)\n# 100 loops, best of 5: 6.82 ms per loop\n%timeit func2_nb(f_arr, nech, n)\n# 100 loops, best of 5: 3.05 ms per loop\n%timeit func_pdisty(f_arr, nech, n)\n# 1 loop, best of 5: 344 ms per loop\n%timeit func2_cy(f_arr, nech, n)\n# 100 loops, best of 5: 3.95 ms per loop\n
Run Code Online (Sandbox Code Playgroud)\n

...以及您提供的输入大小:

\n
nech = n = 4000\nf_arr = np.arange(n)\n%timeit func_OP(f_arr, nech, n)\n# 1 loop, best of 5: 24.3 s per loop\n%timeit func(f_arr, nech, n)\n# 1 loop, best of 5: 9.27 s per loop\n%timeit func2(f_arr, nech, n)\n# 1 loop, best of 5: 5.71 s per loop\n%timeit func_OP_nb(f_arr, nech, n)\n# 10 loops, best of 5: 29 ms per loop\n%timeit func_nb(f_arr, nech, n)\n# 10 loops, best of 5: 27.3 ms per loop\n%timeit func2_nb(f_arr, nech, n)\n# 100 loops, best of 5: 12.2 ms per loop\n%timeit func_pdisty(f_arr, nech, n)\n# 1 loop, best of 5: 706 ms per loop\n%timeit func2_cy(f_arr, nech, n)\n# 100 loops, best of 5: 15.9 ms per loop\n
Run Code Online (Sandbox Code Playgroud)\n


And*_*eak 6

免责声明:正如 @norok2 在评论中指出的那样,N*(N-1)/2由于使用了pdist. 对于您来说N = 70299,这意味着数组中有大约 18.5 GB 的双精度数。其他索引数组将具有类似的大小。因此,除非您的某些用例具有较小的N,否则此答案中的矢量化方法仅在您拥有大量内存时才可行。

\n
\n

正如其他人所指出的,仅仅将代码从一种语言翻译成另一种语言并不会产生两种语言的最佳代码。单独使用 Cython 并不能保证加速,就像单独使用 NumPy 不能保证加速一样。

\n

Norok2 的答案很好地向您展示了如何使用numba或类似的东西来编译您的数字代码。这可以为您提供与 MATLAB 中的性能非常相似的功能,因为 MATLAB 有自己的即时 (JIT) 编译器。还有优化代码编译的回旋余地,因为多种实现最终可能会带来截然不同的性能。

\n

无论如何,我想指出的是,您还可以通过使用 NumPy 和 SciPy 中的高级功能来加速代码。特别是,您想要计算一组 1d 点之间的成对平方距离。这就是scipy.spatial.distance.pdist可以为您做的(使用\'sqeuclidean\'平方欧几里得范数)。优点是它只计算每个成对距离一次(这对于 CPU 和内存性能来说非常好),但缺点是挑选出你想要总结的贡献有点麻烦。

\n

无论如何,这里是与您的 Python 实现相对应的代码(使用内部循环使用的修复np.arange(1, N-i)而不是np.arange(1, N-i-1)):

\n
from scipy.spatial.distance import pdist\n\ndef pdisty(f, nech):\n    offset = f.size - nech\n    res = np.zeros_like(f, shape=nech)\n    dists = pdist(f[:, None], metric=\'sqeuclidean\')\n    counts = np.arange(offset, f.size)[::-1]\n    inds = np.repeat(np.arange(res.size), counts)\n    np.add.at(res, inds, dists[:inds.size])\n    res /= counts\n    return res\n
Run Code Online (Sandbox Code Playgroud)\n

这里发生的事情是

\n
    \n
  1. 我们计算每对唯一的数组值的成对距离并将其存储到dists
  2. \n
  3. 我们计算每个点涉及的对数(这是我们最后必须标准化的),将其存储到counts
  4. \n
  5. 找出dists它对应的每个值的一维索引(这是困难的部分),将其存储在inds
  6. \n
  7. 用于np.add.at累积对适当产出指数的每个贡献
  8. \n
  9. 用计数标准化。
  10. \n
\n

以下是 的一些计时N = 1000,其中是Norok2 的答案func2()中的相应函数:

\n
>>> %timeit structFunPython(f, f.size - 1)\n... %timeit func2(f, f.size - 1)\n... %timeit pdisty(f, f.size - 1)\n1.48 s \xc2\xb1 89.6 ms per loop (mean \xc2\xb1 std. dev. of 7 runs, 1 loop each)\n274 ms \xc2\xb1 2.71 ms per loop (mean \xc2\xb1 std. dev. of 7 runs, 1 loop each)\n36.7 ms \xc2\xb1 1.21 ms per loop (mean \xc2\xb1 std. dev. of 7 runs, 10 loops each)\n
Run Code Online (Sandbox Code Playgroud)\n

上面的解决方案要快得多,但当然它仍然比完全编译的解决方案慢。如果您遇到其他依赖项或在系统上安装 llvm 的问题,这可能是一个合理的折衷方案。最重要的是,代码应该适应您尝试优化它的语言。

\n
\n

为了完整起见,这里是我用于比较的实现(我稍微更改了签名,因为N可以从输入数组计算,并且我修复了一些栅栏错误):

\n
>>> %timeit structFunPython(f, f.size - 1)\n... %timeit func2(f, f.size - 1)\n... %timeit pdisty(f, f.size - 1)\n1.48 s \xc2\xb1 89.6 ms per loop (mean \xc2\xb1 std. dev. of 7 runs, 1 loop each)\n274 ms \xc2\xb1 2.71 ms per loop (mean \xc2\xb1 std. dev. of 7 runs, 1 loop each)\n36.7 ms \xc2\xb1 1.21 ms per loop (mean \xc2\xb1 std. dev. of 7 runs, 10 loops each)\n
Run Code Online (Sandbox Code Playgroud)\n

通过这些定义,所有三个函数在机器精度内给出相同的结果:

\n
def structFunPython(f, nech):\n    """Very slightly modified from the question"""\n    N = f.size\n    sf2 = np.zeros(nech)\n    count = np.zeros(nech)\n    for i in range(nech):\n        indN = np.arange(1,N-i)\n        for j in indN:\n            sf2[i] += np.power((f[i+j]-f[i]),2)\n            count[i] += 1\n        sf2[i] = sf2[i]/count[i]\n    return sf2\n\n\ndef func2(f_arr, nech):\n    """Very slightly modified from norok2\'s answer\n\n    See https://stackoverflow.com/a/71704834/5067311\n\n    """\n    n = f_arr.size\n    sf2 = np.zeros(nech)\n    for i in range(nech):\n        for j in range(1, n - i):\n            sf2[i] += (f_arr[i + j] - f_arr[i]) ** 2\n        sf2[i] /= (n - i - 1)\n    return sf2\n
Run Code Online (Sandbox Code Playgroud)\n


jon*_*oni 5

您已经得到了一些很好的答案,这些答案提供了快速的 numpy 和 numba 实现。我想提供一个不错的 Cython 实现以进行公平比较。

\n

首先,让我们在我的机器上对 Norok2 最快的 numba 实现进行计时:

\n
In [3]: %timeit func2_nb(f_arr, n)\n3.4 s \xc2\xb1 129 ms per loop (mean \xc2\xb1 std. dev. of 7 runs, 1 loop each)\n
Run Code Online (Sandbox Code Playgroud)\n
    \n
  • 为了在 Cython 中获得快速内存访问,向其提供有关输入数组的所有必要信息至关重要。在您的情况下,f_arr是一个 C 连续的 np.ndarray ,因此我们使用 C 连续的内存视图double[::1]而不是普通的内存视图double[:]。不同之处在于,索引 C 连续内存视图会减少为纯 C 代码f_arr[i],而后者则减少为f_arr[i + f_arr.strides[0]].

    \n
  • \n
  • 接下来,值得一提的是,Python 的幂运算符a**2将被 C 代码取代pow(a,2),即我们称之为pow-function。即使在 C 语言中,在紧密循环中调用函数也会产生不必要的开销。在 C 语言中,我们只需编写a*a. 那么让我们在 Cython 中做同样的事情。

    \n
  • \n
\n

在代码中:

\n
%%cython -c=-O3 -c=-march=native\n# for MSVC use: -c=/Ox -c=/arch:AVX2\n#           or: -c=/Ox -c=/arch:AVX512 if your CPU supports AVX512 \n\nimport numpy as np\ncimport numpy as np\ncimport cython\n\n@cython.boundscheck(False)\n@cython.wraparound(False)\ndef func2_cy(double[::1] f_arr, int nech):\n    cdef int n = f_arr.size\n    cdef double[::1] sf2 = np.zeros(nech)\n    cdef int i, j\n    for i in range(nech):\n        for j in range(n-i-2):\n            sf2[i] += (f_arr[i+j+1]-f_arr[i])*(f_arr[i+j+1]-f_arr[i])\n        sf2[i] /= (n - i - 2)\n    return np.asarray(sf2)\n
Run Code Online (Sandbox Code Playgroud)\n

在我的机器上使用 macOS 上的 Apple Clang 13.1.6 进行编译,这比前面提到的 numba 实现要慢:

\n
In [5]: %timeit func2_cy(f_arr, n)\n4.2 s \xc2\xb1 32 ms per loop (mean \xc2\xb1 std. dev. of 7 runs, 1 loop each)\n
Run Code Online (Sandbox Code Playgroud)\n

然而,人们应该意识到 Cython 基本上只是一个Python 到 C 的编译器。这意味着,在删除所有 Python 交互后,我们可以尝试使用 C 代码时应用的相同性能技巧:传递优化标志-O3并启用当前平台上可用的所有 CPU 指令-march=native。请注意,这也意味着良好的 Cython 代码(即在紧密循环中没有 Python 交互的代码)的性能在很大程度上取决于您的 C 编译器。根据我的经验,由于 MSVC 的自动矢量化效果不佳,numba 通常比 Windows 上的 Cython 更快。在 macOS/Linux 和 gcc 或 clang 上,情况通常是相反的。

\n

这些众所周知的性能技巧之一是循环展开,以便为编译器提供 SIMD 向量化特定循环的提示。展开最里面的循环,函数如下所示:

\n
%%cython -c=-O3 -c=-march=native\n# for MSVC use: -c=/Ox -c=/arch:AVX2\n#           or: -c=/Ox -c=/arch:AVX512 if your CPU supports AVX512 \n\n#cython: cdivision=True\n\nimport numpy as np\ncimport numpy as np\ncimport cython\n\n@cython.boundscheck(False)\n@cython.wraparound(False)\ndef func3_cy(double[::1] f_arr, int nech):\n    cdef int n = f_arr.size\n    cdef double[::1] sf2 = np.zeros(nech)\n    cdef double fi, fij, fij1, fij2, fij3\n    cdef int i, j, ntmp\n    for i in range(nech):\n        # unroll the inner loop for better SIMD vectorization\n        j = 1\n        ntmp = (n-i-2) - ((n-i-2) % 4)\n        # we use a while-loop since Cython doesn\'t support range(j, ntmp, 4)\n        while j < ntmp:\n            # unpack the values and hope the CPU keeps them in its register\n            fi   = f_arr[i]\n            fij  = f_arr[i+j+1]\n            fij1 = f_arr[i+j+2]\n            fij2 = f_arr[i+j+3]\n            fij3 = f_arr[i+j+4]\n            sf2[i] += ((fij-fi)*(fij-fi) + (fij1-fi)*(fij1-fi) + (fij2-fi)*(fij2-fi) + (fij3-fi)*(fij3-fi))\n            j += 4\n        for j in range(ntmp, n - i - 2):\n            sf2[i] += (f_arr[i+j+1] - f_arr[i]) * (f_arr[i+j+1] - f_arr[i])\n        sf2[i] /= (n-i-2)\n    return np.asarray(sf2)\n
Run Code Online (Sandbox Code Playgroud)\n

在我的机器上,这几乎比 numba 快两倍:

\n
In [7]: %timeit func3_cy(f_arr, n)\n1.7 s \xc2\xb1 54.4 ms per loop (mean \xc2\xb1 std. dev. of 7 runs, 1 loop each)\n
Run Code Online (Sandbox Code Playgroud)\n

更进一步,我们可以借助OpenMPprange的一个薄包装来并行化外循环:

\n
%%cython -c=-O3 -c=-march=native -c=-fopenmp --link-args=-fopenmp\n# for MSVC use: -c=/Ox -c=/openmp -c=/arch:AVX2\n#           or: -c=/Ox -c=/openmp -c=/arch:AVX512 if your CPU supports AVX512 \n\n#cython: cdivision=True\n\nimport numpy as np\ncimport numpy as np\ncimport cython\nfrom cython.parallel cimport prange\n\n@cython.boundscheck(False)\n@cython.wraparound(False)\ndef func4_cy(double[::1] f_arr, int nech):\n    cdef int n = f_arr.size\n    cdef double[::1] sf2 = np.zeros(nech)\n    cdef double fi, fij, fij1, fij2, fij3\n    cdef int i, j, ntmp\n    for i in prange(nech, nogil=True, schedule="static", chunksize=1):\n        # unroll the inner loop for better SIMD vectorization\n        j = 1\n        ntmp = (n-i-2) - ((n-i-2) % 4)\n        # we use a while-loop since Cython doesn\'t support range(j, ntmp, 4)\n        while j < ntmp:\n            # unpack the values and hope the CPU keeps them in its register\n            fi   = f_arr[i]\n            fij  = f_arr[i+j+1]\n            fij1 = f_arr[i+j+2]\n            fij2 = f_arr[i+j+3]\n            fij3 = f_arr[i+j+4]\n            sf2[i] += ((fij-fi)*(fij-fi) + (fij1-fi)*(fij1-fi) + (fij2-fi)*(fij2-fi) + (fij3-fi)*(fij3-fi))\n            j += 4\n        for j in range(ntmp, n - i - 2):\n            sf2[i] += (f_arr[i+j+1] - f_arr[i]) * (f_arr[i+j+1] - f_arr[i])\n        sf2[i] /= (n-i-2)\n    return np.asarray(sf2)\n
Run Code Online (Sandbox Code Playgroud)\n

在我的 8 个 CPU 核心的机器上,这是迄今为止最快的实现:

\n
In [9]: %timeit func4_cy(f_arr, n)\n329 ms \xc2\xb1 15.5 ms per loop (mean \xc2\xb1 std. dev. of 7 runs, 1 loop each)\n
Run Code Online (Sandbox Code Playgroud)\n

不过,值得一提的是,Numba 还支持线程并行性,因此我期望 Numba 具有类似的性能。

\n

  • @AndrasDeak--СлаваУкраїні 至少现在它应该给出与OP的python版本相同的结果。然而,我只是注意到OP的Matlab版本计算`(f(i+j)-f(j))**2`,而他的Python实现计算`(f[i+j]-f[i])**2 `.. 请注意,它是 `f[i]` 而不是 `f[j]`。 (2认同)