我有一个分析代码，它使用 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.- 这不是我需要进行的计算，只是一个显示相同内容的简单示例。

最佳答案

稍作修改，版本 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。

关于python - Numpy vs Cython 速度，我们在Stack Overflow上找到一个类似的问题： https://stackoverflow.com/questions/7799977/