Numpy与Cython速度比较

Question

我有一段分析代码，它使用numpy进行一些复杂的数字运算。出于好奇，我尝试用cython编译它，只做了很少的修改，然后我把numpy的部分改成了用循环来写。

让我惊讶的是，基于循环的代码运行得快得多（快了8倍）。我不能发布完整的代码，但我整理了一个非常简单的无关计算，显示了类似的情况（虽然时间差别没有那么大）：

版本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会更快呢？

补充说明：这不是我需要进行的计算，只是一个简单的例子，展示了相同的情况。

Answer 1

我建议使用 -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 实际上比其他两个版本稍慢一点。

Cython 在处理 numpy 无法以“向量化”方式进行的数组操作时表现得特别好，或者在进行一些内存密集型的操作时，它可以帮助你避免创建一个大的临时数组。我在自己的一些代码中，使用 Cython 相比 numpy 提升了 115 倍的速度：

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

其中一部分是直接在 C 代码层面调用 randomkit，而不是通过 numpy.random 来调用，但大部分的提升是因为 Cython 将计算密集型的 for 循环翻译成了纯 C 代码，而没有调用 Python。

Answer 2

稍微修改一下，版本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 3

正如其他回答提到的，版本2和版本1基本上是一样的，因为cython无法深入优化数组访问操作。造成这种情况的原因有两个：

• 首先，每次调用numpy函数时都会有一些额外的开销，这和优化过的C代码相比是有差距的。不过，如果每次操作处理的是大数组，这种开销就不那么重要了。

• 其次，会创建中间数组。我们来看一个更复杂的操作，比如 out[row, :] = A[row, :] + B[row, :]*C[row, :] 。在这种情况下，必须在内存中创建一个完整的数组 B*C，然后再加到 A 上。这就意味着CPU缓存会被频繁使用，因为数据需要从内存中读取和写入，而不是直接在CPU中使用。特别是当你处理大数组时，这个问题会更加严重。

特别是你提到你的真实代码比示例更复杂，并且显示出更大的加速效果，我怀疑第二个原因可能是你情况中的主要因素。

顺便说一下，如果你的计算足够简单，你可以通过使用 numexpr 来克服这个问题，当然cython在很多其他情况下也很有用，所以可能对你来说是更好的选择。