使用Cython加速数组广播中的循环

Question

总结：

你们真是太棒了……我的真实代码终于能正常工作了。我听取了JoshAdel的建议，具体包括：

1) 把所有的ndarray改成了类型化的内存视图
2) 手动展开了所有的numpy数组计算
3) 使用静态定义的无符号整数作为索引
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

不过，现在我有这种代码，它不是单个数字的循环，而是数组的循环。（是的……我在解一个偏微分方程，所以会这样）。

我知道下面这个例子看起来很傻，但只是想让你理解这种计算的类型：

纯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

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)
    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

基本上，你的理解是错的。Python在这类循环上表现得非常好。

这是因为这些循环的开销比较大。正如Josh Adel所展示的，使用纯C语言的变体速度提升也只是2到4倍。

解决这个问题的最好方法是找出几种低效的情况：

• 很多的切片（比如切割很多小行）

  这个问题最好通过手动迭代来解决，并且可以考虑使用Cython或Numba（或者其他即时编译器）

• 完整的Numpy数组计算速度不够快

  Numba和Parakeet可以提供帮助，但你更希望使用Theano或numexpr

• 更复杂的、不能即时编译的慢操作

  在这种情况下，Cython或C是你唯一的选择，所以要转向它们

也可以看看这个回答的第一部分。

Answer 2

在这两个其他的实现中，我尝试了一些方法：

• 使用cython编译器指令来去掉一些numpy通常需要进行的检查
• 使用类型化的内存视图，这样我可以指定内存的布局（有时候它们的速度比旧的缓冲区接口更快）
• 展开循环，这样我们就不使用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毫秒
• 你原来的带类型的cython函数 i：1次循环，3次中最好的一次：每次循环428毫秒
• func2：1次循环，3次中最好的一次：每次循环336毫秒
• func3：1次循环，3次中最好的一次：每次循环206毫秒