Cython 通过阵列广播加速循环

Question

总结：

你们太棒了...我的真实代码工作正常。我采纳了 JoshAdel 的建议，即：

1) 将所有 ndarray 更改为类型化内存视图 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，所以会发生这种情况）。

我知道下面的例子很愚蠢，但只是尝试了解计算类型：

纯蟒蛇：

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 通常必须进行的一些检查
• 使用类型化的内存视图，以便我可以指定内存布局（有时它们通常比旧的缓冲区接口更快）
• 展开循环以便我们不使用 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 毫秒

Answer 2

很大程度上，你错了。 Python 在这些类型的循环中表现出色。

这是因为它们是重量级的。正如 Josh Adel 所示，一个有效的纯 C 变体只能实现 2 到 4 倍的加速。

最好的办法是寻找几种类型的低效率：

• 很多个切片（比如切片很多小行）

  最好通过手动迭代并移至 Cython 或 Numba（或其他 JITing 编译器）来解决此问题

• 不够快的全面 Numpy 数组计算

  Numba 和 Parakeet 可以 提供帮助，但您确实需要 Theano 或 numexpr

• 更复杂、非 JITable、不可表达的慢东西

  Cython 或 C 是您唯一的选择，所以转向他们

See also the first part of this answer.