使用 NumPy

注意

Cython 0.16 引入了类型化内存视图,作为此处描述的 NumPy 集成的继任者。它们比下面的缓冲区语法更容易使用,开销更低,并且可以在不使用 GIL 的情况下传递。它们应该优先于此页面中介绍的语法。请参阅 适用于 NumPy 用户的 Cython

注意

目前无法使用 Python 风格的注释来有效地指定 Numpy 数组,我们目前也不打算添加。如果您想使用注释类型,我们建议使用类型化内存视图。

您可以在 Cython 中使用 NumPy,就像在普通 Python 中一样,但这样做会失去潜在的高速提升,因为 Cython 支持快速访问 NumPy 数组。让我们看看它如何在简单的示例中工作。

下面的代码对图像执行 2D 离散卷积,并使用一个滤波器(我相信你可以做得更好!,让它作为演示目的)。它既是有效的 Python 代码,也是有效的 Cython 代码。我将它称为 convolve_py.py(用于 Python 版本)和 convolve1.pyx(用于 Cython 版本)——Cython 使用“.pyx”作为其文件后缀。

import numpy as np


def naive_convolve(f, g):
    # f is an image and is indexed by (v, w)
    # g is a filter kernel and is indexed by (s, t),
    #   it needs odd dimensions
    # h is the output image and is indexed by (x, y),
    #   it is not cropped
    if g.shape[0] % 2 != 1 or g.shape[1] % 2 != 1:
        raise ValueError("Only odd dimensions on filter supported")
    # smid and tmid are number of pixels between the center pixel
    # and the edge, ie for a 5x5 filter they will be 2.
    #
    # The output size is calculated by adding smid, tmid to each
    # side of the dimensions of the input image.
    vmax = f.shape[0]
    wmax = f.shape[1]
    smax = g.shape[0]
    tmax = g.shape[1]
    smid = smax // 2
    tmid = tmax // 2
    xmax = vmax + 2 * smid
    ymax = wmax + 2 * tmid
    # Allocate result image.
    h = np.zeros([xmax, ymax], dtype=f.dtype)
    # Do convolution
    for x in range(xmax):
        for y in range(ymax):
            # Calculate pixel value for h at (x,y). Sum one component
            # for each pixel (s, t) of the filter g.
            s_from = max(smid - x, -smid)
            s_to = min((xmax - x) - smid, smid + 1)
            t_from = max(tmid - y, -tmid)
            t_to = min((ymax - y) - tmid, tmid + 1)
            value = 0
            for s in range(s_from, s_to):
                for t in range(t_from, t_to):
                    v = x - smid + s
                    w = y - tmid + t
                    value += g[smid - s, tmid - t] * f[v, w]
            h[x, y] = value
    return h

这应该编译成 yourmod.so(对于 Linux 系统,在 Windows 系统上,它将是 yourmod.pyd)。我们运行一个 Python 会话来测试 Python 版本(从 .py 文件导入)和编译后的 Cython 模块。

In [1]: import numpy as np
In [2]: import convolve_py
In [3]: convolve_py.naive_convolve(np.array([[1, 1, 1]], dtype=np.int64),
...     np.array([[1],[2],[1]], dtype=np.int64))
Out [3]:
array([[1, 1, 1],
    [2, 2, 2],
    [1, 1, 1]])
In [4]: import convolve1
In [4]: convolve1.naive_convolve(np.array([[1, 1, 1]], dtype=np.int64),
...     np.array([[1],[2],[1]], dtype=np.int64))
Out [4]:
array([[1, 1, 1],
    [2, 2, 2],
    [1, 1, 1]])
In [11]: N = 100
In [12]: f = np.arange(N*N, dtype=np.int64).reshape((N,N))
In [13]: g = np.arange(81, dtype=np.int64).reshape((9, 9))
In [19]: %timeit -n2 -r3 convolve_py.naive_convolve(f, g)
2 loops, best of 3: 1.86 s per loop
In [20]: %timeit -n2 -r3 convolve1.naive_convolve(f, g)
2 loops, best of 3: 1.41 s per loop

目前还没有那么大的区别;因为 C 代码仍然完全执行 Python 解释器所做的操作(例如,为使用的每个数字分配一个新对象)。查看生成的 html 文件,看看即使是最简单的语句都需要什么,你就能很快明白。我们需要给 Cython 提供更多信息;我们需要添加类型。

添加类型

要添加类型,我们使用自定义 Cython 语法,因此我们现在正在破坏 Python 源代码兼容性。考虑以下代码(*阅读注释!*)

# tag: numpy
# You can ignore the previous line.
# It's for internal testing of the cython documentation.

import numpy as np

# "cimport" is used to import special compile-time information
# about the numpy module (this is stored in a file numpy.pxd which is
# distributed with Numpy).
# Here we've used the name "cnp" to make it easier to understand what
# comes from the cimported module and what comes from the imported module,
# however you can use the same name for both if you wish.
cimport numpy as cnp

# It's necessary to call "import_array" if you use any part of the
# numpy PyArray_* API. From Cython 3, accessing attributes like
# ".shape" on a typed Numpy array use this API. Therefore we recommend
# always calling "import_array" whenever you "cimport numpy"
cnp.import_array()

# We now need to fix a datatype for our arrays. I've used the variable
# DTYPE for this, which is assigned to the usual NumPy runtime
# type info object.
DTYPE = np.int64

# "ctypedef" assigns a corresponding compile-time type to DTYPE_t. For
# every type in the numpy module there's a corresponding compile-time
# type with a _t-suffix.
ctypedef cnp.int64_t DTYPE_t

# "def" can type its arguments but not have a return type. The type of the
# arguments for a "def" function is checked at run-time when entering the
# function.
#
# The arrays f, g and h is typed as "np.ndarray" instances. The only effect
# this has is to a) insert checks that the function arguments really are
# NumPy arrays, and b) make some attribute access like f.shape[0] much
# more efficient. (In this example this doesn't matter though.)
def naive_convolve(cnp.ndarray f, cnp.ndarray g):
    if g.shape[0] % 2 != 1 or g.shape[1] % 2 != 1:
        raise ValueError("Only odd dimensions on filter supported")
    assert f.dtype == DTYPE and g.dtype == DTYPE

    # The "cdef" keyword is also used within functions to type variables. It
    # can only be used at the top indentation level (there are non-trivial
    # problems with allowing them in other places, though we'd love to see
    # good and thought out proposals for it).
    #
    # For the indices, the "int" type is used. This corresponds to a C int,
    # other C types (like "unsigned int") could have been used instead.
    # Purists could use "Py_ssize_t" which is the proper Python type for
    # array indices.
    cdef int vmax = f.shape[0]
    cdef int wmax = f.shape[1]
    cdef int smax = g.shape[0]
    cdef int tmax = g.shape[1]
    cdef int smid = smax // 2
    cdef int tmid = tmax // 2
    cdef int xmax = vmax + 2 * smid
    cdef int ymax = wmax + 2 * tmid
    cdef cnp.ndarray h = np.zeros([xmax, ymax], dtype=DTYPE)
    cdef int x, y, s, t, v, w

    # It is very important to type ALL your variables. You do not get any
    # warnings if not, only much slower code (they are implicitly typed as
    # Python objects).
    cdef int s_from, s_to, t_from, t_to

    # For the value variable, we want to use the same data type as is
    # stored in the array, so we use "DTYPE_t" as defined above.
    # NB! An important side-effect of this is that if "value" overflows its
    # datatype size, it will simply wrap around like in C, rather than raise
    # an error like in Python.
    cdef DTYPE_t value
    for x in range(xmax):
        for y in range(ymax):
            s_from = max(smid - x, -smid)
            s_to = min((xmax - x) - smid, smid + 1)
            t_from = max(tmid - y, -tmid)
            t_to = min((ymax - y) - tmid, tmid + 1)
            value = 0
            for s in range(s_from, s_to):
                for t in range(t_from, t_to):
                    v = x - smid + s
                    w = y - tmid + t
                    value += g[smid - s, tmid - t] * f[v, w]
            h[x, y] = value
    return h

在构建完此代码并继续我的(非常非正式的)基准测试后,我得到

In [21]: import convolve2
In [22]: %timeit -n2 -r3 convolve2.naive_convolve(f, g)
2 loops, best of 3: 828 ms per loop

高效索引

仍然有一个瓶颈影响性能,那就是数组查找和赋值。 [] 运算符仍然使用完整的 Python 操作——我们想要做的是以 C 速度直接访问数据缓冲区。

我们需要做的就是对 ndarray 对象的内容进行类型化。我们使用特殊的“缓冲区”语法来完成此操作,该语法必须告知数据类型(第一个参数)和维度数(“ndim”关键字参数,如果未提供,则假定为一维)。

这些是必要的更改

...
def naive_convolve(cnp.ndarray[DTYPE_t, ndim=2] f, cnp.ndarray[DTYPE_t, ndim=2] g):
...
cdef cnp.ndarray[DTYPE_t, ndim=2] h = ...

用法

In [18]: import convolve3
In [19]: %timeit -n3 -r100 convolve3.naive_convolve(f, g)
3 loops, best of 100: 11.6 ms per loop

请注意此更改的重要性。

陷阱:这种高效索引只影响某些索引操作,即那些具有完全 ndim 个类型化整数索引的操作。因此,如果 v 没有类型化,那么查找 f[v, w] 不会被优化。另一方面,这意味着您可以继续使用 Python 对象进行复杂的动态切片等操作,就像数组没有类型化一样。

进一步调整索引

数组查找仍然受到两个因素的减速

  1. 执行边界检查。

  2. 负索引被检查并正确处理。上面的代码被明确地编码,以便它不使用负索引,并且它(希望)始终在边界内访问。我们可以添加一个装饰器来禁用边界检查

    ...
    cimport cython
    @cython.boundscheck(False) # turn off bounds-checking for entire function
    @cython.wraparound(False)  # turn off negative index wrapping for entire function
    def naive_convolve(cnp.ndarray[DTYPE_t, ndim=2] f, cnp.ndarray[DTYPE_t, ndim=2] g):
    ...
    

现在没有执行边界检查(作为副作用,如果你碰巧访问了边界之外,最好的情况是你的程序会崩溃,最坏的情况是数据会损坏)。可以通过多种方式切换边界检查模式,有关更多信息,请参阅编译器指令

此外,我们还禁用了检查负索引的包装(例如,g[-1] 给出最后一个值)。与禁用边界检查一样,如果我们尝试在禁用此功能的情况下实际使用负索引,将会发生不好的事情。

函数调用开销现在开始发挥作用,因此我们将后两个示例与更大的 N 进行比较。

In [11]: %timeit -n3 -r100 convolve4.naive_convolve(f, g)
3 loops, best of 100: 5.97 ms per loop
In [12]: N = 1000
In [13]: f = np.arange(N*N, dtype=np.int64).reshape((N,N))
In [14]: g = np.arange(81, dtype=np.int64).reshape((9, 9))
In [17]: %timeit -n1 -r10 convolve3.naive_convolve(f, g)
1 loops, best of 10: 1.16 s per loop
In [18]: %timeit -n1 -r10 convolve4.naive_convolve(f, g)
1 loops, best of 10: 597 ms per loop

(这也是一个混合基准测试,因为结果数组是在函数调用中分配的。)

警告

速度是有代价的。特别是将类型化对象(如示例代码中的 fgh)设置为 None 可能很危险。将此类对象设置为 None 是完全合法的,但你只能做的事情是检查它们是否为 None。所有其他使用(属性查找或索引)都可能导致段错误或数据损坏(而不是像在 Python 中那样引发异常)。

实际规则要复杂一些,但主要信息很清楚:不要在不知道它们没有设置为 None 的情况下使用类型化对象。

类型化不做什么

将事物类型化为 ndarray 的主要目的是允许对单个元素进行高效索引,并加快对少量属性(如 .shape)的访问速度。类型化不允许多个 Cython 加快对整个数组的数学运算(例如,将两个数组加在一起)。类型化不允许多个 Cython 加快对 Numpy 全局函数或数组方法的调用。

更通用的代码

可以这样做

def naive_convolve(object[DTYPE_t, ndim=2] f, ...):

即使用 object 而不是 cnp.ndarray。在 Python 3.0 下,这可以允许你的算法使用任何支持缓冲区接口的库;并且对例如 Python 图像库的支持可以很容易地添加,如果有人也对 Python 2.x 感兴趣。

不过,这样做会有一些速度损失(因为如果类型设置为 cnp.ndarray,你会在编译时做出更多假设,具体来说,假设数据存储在纯步长模式下,而不是间接模式下)。

缓冲区选项

创建缓冲区类型时,接受以下选项

  • ndim - 维数的整数。

  • mode - 来自以下字符串

    • "c" - C 连续数组,

    • "fortran" - Fortran 连续数组,

    • "strided" - 对单个内存块的非连续查找,

    • "full" - 任何有效的缓冲区,包括间接数组。

  • negative_indices - 布尔值,指定是否允许负索引,本质上是编译器指令 cython.wraparound 的每个变量版本。

  • cast - 布尔值,指定是否允许用户将数组视为不同类型。源类型和目标类型的大小必须相同。在 C++ 中,这等效于 reinterpret_cast

在所有情况下,这些参数都必须是编译时常量。

作为指定参数的示例

cdef cnp.ndarray[double, ndim=2, mode="c", cast=True] some_array

cast 可用于获取具有非原生字节序的数组的低级视图

cdef cnp.ndarray[cnp.uint32, cast=True] values = np.arange(10, dtype='>i4')

虽然正确解释转换后的数据是用户的责任。