python - Cython 通过数组广播加速循环

Question

总结:

你们太棒了……我的真实代码可以运行了。我采纳了 JoshAdel 的建议，即:

1) 将所有 ndarray 更改为类型化的内存 View 2) 手动展开所有 numpy 数组计算 3) 为索引使用静态定义的 unsigned int 4) 禁用边界检查和环绕

另外，非常感谢 Veedrac 的洞察力!

原帖:

我知道 python 执行这类代码真的很慢:

import numpy as np

def func0():
    x = 0.
    for i in range(1000):
        x += 1.
    return

如果我将其更改为 Cython，它会更快:

import numpy as np
cimport numpy as np

def func0():
    cdef double x = 0.
    for i in range(1000):
        x += 1.
    return

结果如下:

# Python
10000 loops, best of 3: 58.9 µs per loop
# Cython
10000000 loops, best of 3: 66.8 ns per loop

但是，现在我有这种代码，它不是单个数字的循环，而是数组的循环。 (是的...我正在求解 PDE，所以会发生这种情况)。

我知道下面的例子很愚蠢，但只是想了解一下计算类型:

纯 python :

def func1():
    array1 = np.random.rand(50000, 4)
    array2 = np.random.rand(50000)
    for i in range(1000):
        array1[:, 0] += array2
        array1[:, 1] += array2
        array1[:, 2] += array2
        array1[:, 3] += array2
    return

赛通:

def func1():
    cdef np.ndarray[np.double_t, ndim=2] array1 = np.random.rand(50000, 4)
    cdef np.ndarray[np.double_t, ndim=1] array2 = np.random.rand(50000)
    for i in range(1000):
        array1[:, 0] += array2
        array1[:, 1] += array2
        array1[:, 2] += array2
        array1[:, 3] += array2
    return

而且几乎没有任何改善。与此同时，我知道 Python 不擅长处理这些巨大的循环，因为开销很大。

# Python
1 loops, best of 3: 299 ms per loop
# Cython
1 loops, best of 3: 300 ms per loop

我应该如何改进这类代码有什么建议吗？

Answer 1

在这两个其他实现中，我尝试过

• 使用 cython 编译器指令删除 numpy 通常必须进行的一些检查
• 使用类型化的内存 View ，以便我可以指定内存布局(有时它们通常比旧的缓冲区接口(interface)更快)
• 展开循环以便我们不使用 numpy 的切片机制:

否则，您只是通过 cython 使用 numpy，而 numpy 已经在后台以相当快的 c 代码实现了这一点。

方法:

import numpy as np
cimport numpy as np

cimport cython

def func1():
    cdef np.ndarray[np.double_t, ndim=2] array1 = np.random.rand(50000, 4)
    cdef np.ndarray[np.double_t, ndim=1] array2 = np.random.rand(50000)
    cdef unsigned int i
    for i in range(1000):
        array1[:, 0] += array2
        array1[:, 1] += array2
        array1[:, 2] += array2
        array1[:, 3] += array2
    return

@cython.boundscheck(False)
@cython.wraparound(False)
def func2():
    cdef np.ndarray[np.double_t, ndim=2] array1 = np.random.rand(50000, 4)
    cdef np.ndarray[np.double_t, ndim=1] array2 = np.random.rand(50000)
    cdef unsigned int i, k
    for i in range(1000):
        for k in xrange(50000):
            array1[k, 0] += array2[k]
            array1[k, 1] += array2[k]
            array1[k, 2] += array2[k]
            array1[k, 3] += array2[k]
    return


@cython.boundscheck(False)
@cython.wraparound(False)
def func3():
    cdef np.ndarray[np.double_t, ndim=2] array1 = np.random.rand(50000, 4)
    cdef np.ndarray[np.double_t, ndim=1] array2 = np.random.rand(50000)
    cdef np.double_t[::1] a2 = array2
    cdef np.double_t[:,::1] a1 = array1
    cdef unsigned int i, k

    for i in range(1000):
        for k in xrange(50000):
            a1[k, 0] += a2[k]
            a1[k, 1] += a2[k]
            a1[k, 2] += a2[k]
            a1[k, 3] += a2[k]
    return

计时(在我的机器上——编译器和硬件当然会影响计时):

• 纯 numpy:1 个循环，3 个循环中的最佳:每个循环 419 毫秒
• 输入 i 的原始 cython 函数:1 个循环，3 个循环中的最佳:每个循环 428 毫秒
• func2:1 个循环，最好的 3 个循环:每个循环 336 毫秒
• func3:1 个循环，最好的 3 个循环:每个循环 206 毫秒