Numpy vs Cython 速度

Question

我有一个分析代码，它使用 numpy 执行一些繁重的数值运算。只是为了好奇，尝试用 cython 编译它，只做很少的改动，然后我用循环重写了它的 numpy 部分。

令我惊讶的是，基于循环的代码要快得多 (8x)。我无法发布完整的代码，但我整理了一个非常简单的不相关计算，显示出类似的行为（尽管时间差异不是那么大）：

版本 1（无 cython）

import numpy as np

def _process(array):

    rows = array.shape[0]
    cols = array.shape[1]

    out = np.zeros((rows, cols))

    for row in range(0, rows):
        out[row, :] = np.sum(array - array[row, :], axis=0)

    return out

def main():
    data = np.load('data.npy')
    out = _process(data)
    np.save('vianumpy.npy', out)

版本 2（使用 cython 构建模块）

import cython
cimport cython

import numpy as np
cimport numpy as np

DTYPE = np.float64
ctypedef np.float64_t DTYPE_t

@cython.boundscheck(False)
@cython.wraparound(False)
@cython.nonecheck(False)
cdef _process(np.ndarray[DTYPE_t, ndim=2] array):

    cdef unsigned int rows = array.shape[0]
    cdef unsigned int cols = array.shape[1]
    cdef unsigned int row
    cdef np.ndarray[DTYPE_t, ndim=2] out = np.zeros((rows, cols))

    for row in range(0, rows):
        out[row, :] = np.sum(array - array[row, :], axis=0)

    return out

def main():
    cdef np.ndarray[DTYPE_t, ndim=2] data
    cdef np.ndarray[DTYPE_t, ndim=2] out
    data = np.load('data.npy')
    out = _process(data)
    np.save('viacynpy.npy', out)

版本 3（使用 cython 构建模块）

import cython
cimport cython

import numpy as np
cimport numpy as np

DTYPE = np.float64
ctypedef np.float64_t DTYPE_t

@cython.boundscheck(False)
@cython.wraparound(False)
@cython.nonecheck(False)
cdef _process(np.ndarray[DTYPE_t, ndim=2] array):

    cdef unsigned int rows = array.shape[0]
    cdef unsigned int cols = array.shape[1]
    cdef unsigned int row
    cdef np.ndarray[DTYPE_t, ndim=2] out = np.zeros((rows, cols))

    for row in range(0, rows):
        for col in range(0, cols):
            for row2 in range(0, rows):
                out[row, col] += array[row2, col] - array[row, col]

    return out

def main():
    cdef np.ndarray[DTYPE_t, ndim=2] data
    cdef np.ndarray[DTYPE_t, ndim=2] out
    data = np.load('data.npy')
    out = _process(data)
    np.save('vialoop.npy', out)

使用 data.npy 中保存的 10000x10 矩阵，时间为：

$ python -m timeit -c "from version1 import main;main()"
10 loops, best of 3: 4.56 sec per loop

$ python -m timeit -c "from version2 import main;main()"
10 loops, best of 3: 4.57 sec per loop

$ python -m timeit -c "from version3 import main;main()"
10 loops, best of 3: 2.96 sec per loop

这是预期的还是我缺少的优化？版本 1 和 2 给出相同结果的事实在某种程度上是意料之中的，但为什么版本 3 更快？

Ps.- 这不是我需要进行的计算，只是一个显示相同内容的简单示例。

Answer 1

稍作修改，版本 3 的速度提高了一倍：

@cython.boundscheck(False)
@cython.wraparound(False)
@cython.nonecheck(False)
def process2(np.ndarray[DTYPE_t, ndim=2] array):

    cdef unsigned int rows = array.shape[0]
    cdef unsigned int cols = array.shape[1]
    cdef unsigned int row, col, row2
    cdef np.ndarray[DTYPE_t, ndim=2] out = np.empty((rows, cols))

    for row in range(rows):
        for row2 in range(rows):
            for col in range(cols):
                out[row, col] += array[row2, col] - array[row, col]

    return out

计算中的瓶颈是内存访问。您的输入数组是 C 排序的，这意味着沿最后一个轴移动会使内存中的跳跃最小。因此，您的内部循环应沿轴 1，而不是轴 0。进行此更改可将运行时间缩短一半。

如果您需要在小型输入数组上使用此函数，则可以通过使用np.empty 而不是np.ones 来减少开销。为了进一步减少开销，请使用 numpy C API 中的PyArray_EMPTY。

如果您在非常大的输入数组 (2**31) 上使用此函数，则用于索引的整数（以及在 range 函数中）将溢出。为了安全使用：

cdef Py_ssize_t rows = array.shape[0]
cdef Py_ssize_t cols = array.shape[1]
cdef Py_ssize_t row, col, row2

而不是

cdef unsigned int rows = array.shape[0]
cdef unsigned int cols = array.shape[1]
cdef unsigned int row, col, row2

时间：

In [2]: a = np.random.rand(10000, 10)
In [3]: timeit process(a)
1 loops, best of 3: 3.53 s per loop
In [4]: timeit process2(a)
1 loops, best of 3: 1.84 s per loop

process 是您的版本 3。

Answer 2

正如其他答案中提到的，版本 2 与版本 1 基本相同，因为 cython 无法深入研究数组访问运算符以对其进行优化。这有两个原因

• 首先，与优化的 C 代码相比，每次调用 numpy 函数都会产生一定的开销。但是，如果每个操作都处理大型数组，则此开销将变得不那么重要

• 其次，创建中间数组。如果您考虑更复杂的操作，例如out[row, :] = A[row, :] + B[row, :]*C[row, :]，这会更清楚。在这种情况下，整个数组B*C 必须在内存中创建，然后添加到A。这意味着 CPU 缓存正在被破坏，因为数据正在从内存中读取和写入，而不是保存在 CPU 中并立即使用。重要的是，如果您处理大型数组，这个问题会变得更糟。

特别是由于您声明您的实际代码比您的示例更复杂，并且它显示出更大的加速，我怀疑第二个原因可能是您案例的主要因素。

顺便说一句，如果您的计算足够简单，您可以使用numexpr 来克服这种影响，当然 cython 在更多情况下很有用，因此它可能是您更好的方法。

Answer 3

我建议使用 -a 标志让 cython 生成 html 文件，该文件显示正在转换为纯 c 的内容与调用 python API：

http://docs.cython.org/src/quickstart/cythonize.html

版本 2 的结果与版本 1 几乎相同，因为所有繁重的工作都由 Python API（通过 numpy）完成，而 cython 并没有为您做任何事情。事实上在我的机器上，numpy 是针对 MKL 构建的，所以当我使用 gcc 编译 cython 生成的 c 代码时，版本 3 实际上比其他两个慢一点。

当您执行 numpy 无法以“矢量化”方式执行的数组操作时，或者当您执行内存密集型操作时，Cython 会大放异彩，它可以让您避免创建大型临时数组。对于我自己的一些代码，我使用 cython vs numpy 获得了 115 倍的加速：

https://github.com/synapticarbors/pylangevin-integrator

其中一部分是在 c 代码级别调用 randomkit 目录，而不是通过 numpy.random 调用它，但其中大部分是 cython 将计算密集型 for 循环转换为纯 c 而不调用 python。

Answer 4

差异可能是由于版本 1 和 2 对每一行都对 np.sum() 进行 Python 级调用，而版本 3 可能会编译为紧密的纯 C 循环。

研究第 2 版和第 3 版的 Cython 生成的 C 源代码之间的区别应该很有启发性。

Answer 5

我猜您节省的主要开销是创建的临时数组。您创建了一个很大的数组array - array[row, :]，然后使用sum 将其缩减为一个较小的数组。但是构建那个大的临时数组不会是免费的，特别是如果你需要分配内存。