我可以使用 Cython 加速 Numpy 密集型函数的哪些部分

Dav*_*idW 10 python numpy cython

介绍性说明:尝试使用 Cython 加速 Python+Numpy 代码是一个常见问题,这个问题试图创建一个关于可以有效加速哪些类型的操作的规范问题。虽然我试图用一个具体的例子来说明,但这只是为了说明——请不要过多关注这个毫无意义的例子。

另外,我对 Cython 做出了足够的贡献,我应该声明一个从属关系(鉴于我正在提出这个主题)


实际问题

假设我有一个函数尝试对 Numpy 数组进行数值计算。它使用相当典型的操作:

  • 对不易矢量化的数组元素进行循环
  • 调用 Numpy/Scipy 函数(在本例中np.sin)。
  • 对整个数组进行数学运算 ( a-b)
import numpy as np

def some_func(a, b):
    """
    a and b are 1D arrays

    This is intended to be illustrative! Please don't focus on what it
    actually does!
    """
    transformed_a = np.zeros_like(a)
    last = 0
    for n in range(1, a.shape[0]):
        an = a[n]
        if an > 0:
            delta = an - a[n-1]
            transformed_a[n] = delta*last
        else:
            last = np.sin(an)
    return transformed_a * b

a = np.random.randn(100)
b = np.linspace(0, 100, a.shape[0])

print(some_func(a, b))
Run Code Online (Sandbox Code Playgroud)

我可以使用 Cython 加速这个过程吗?我希望哪些部分能够加速?

Dav*_*idW 10

索引各个数组元素

这是 Cython真正可以帮助您的主要代码类型。在 Python 中,对单个元素(例如an = a[n])建立索引可能是一个相当慢的操作。部分原因是 Python 不是一种非常快的语言,因此在循环中多次运行 Python 代码可能会很慢,部分原因是该数组存储为紧凑的 C 浮点数组,但索引操作需要返回 Python目的。因此,索引 Numpy 数组需要分配新的 Python 对象。

在 Cython 中,您可以将数组声明为类型化内存视图np.ndarray. (类型化内存视图是更现代的方法,您通常应该更喜欢它们)。这样做可以让您直接访问紧密封装的 C 数组并检索 C 值,而无需创建 Python 对象。

cython.boundscheck这些指令cython.wraparound对于进一步加快索引速度非常有价值(但请记住它们确实删除了有用的功能,因此在使用它们之前请三思)。

与矢量化

很多时候,Numpy 数组上的循环可以写成向量化操作——一次性作用于整个数组。像这样编写 Python+Numpy 代码通常是个好主意。如果您有多个链式向量化操作,有时值得将其显式编写为 Cython 循环以避免分配中间数组。

或者,鲜为人知的Cython Pythran 后端将一组矢量化 Numpy 运算转换为优化的 C++ 代码。

索引数组切片

在 Cython 中这不是问题,但通常不会单独使您显着加速。

调用 Numpy 函数

例如last = np.sin(an)

这些需要 Python 调用,因此 Cython 通常无法加速这些 - 它无法查看 Numpy 函数的内容。

然而,这里的操作是针对单个值,而不是针对 Numpy 数组。在这种情况下,我们可以使用sinC 标准库,这将比 Python 函数调用快得多。你会from libc.math cimport sin打电话sin而不是np.sin

Numba 是另一种 Python 加速器,它可以更好地了解 Numpy 函数,通常可以在不进行更改的情况下进行优化。

数组分配

例如transformed_a = np.zeros_like(a)

这只是 Numpy 函数调用,因此 Cython 无法加速它。如果它只是一个中间值并且没有返回到 Python 那么你可能会考虑在堆栈上使用固定大小的 C 数组

cdef double transformed_a[10]  # note - you must know the size at compile-time
Run Code Online (Sandbox Code Playgroud)

或者通过 Cmalloc函数分配它们(记住free这一点)。或者使用 Cython 的cython.view.array(它仍然是一个 Python 对象,但可以更快一点)。

全数组算术

例如transformed_a * b,它逐个transformed_a元素b相乘。

Cython 在这里对您没有帮助 - 它只是一个伪装的函数调用(尽管 Pythran+Cython 可能有一些好处)。对于大型数组,这种操作在 Numpy 中非常有效,所以不要想太多。

请注意,整个数组操作没有为 Cython 类型的内存视图定义,因此您需要将np.asarray(memview)它们返回到 Numpy 数组。这通常不需要副本并且速度很快。

对于像这样的某些运算,您可以使用BLASandLAPACK函数(它们是数组和矩阵运算的快速 C 实现)。Scipy 提供了 Cython 接口供他们使用( https://docs.scipy.org/doc/scipy/reference/linalg.cython_blas.html )。它们的使用比自然的 Python 代码稍微复杂一些。

说明性示例

为了完整起见,我会这样写:

import numpy as np
from libc.math cimport sin
cimport cython

@cython.boundscheck(False)
@cython.wraparound(False)
def some_func(double[::1] a, b):
    cdef double[::1] transformed_a = np.zeros_like(a)
    cdef double last = 0
    cdef double an, delta
    cdef Py_ssize_t n
    for n in range(1, a.shape[0]):
        an = a[n]
        if an > 0:
            delta = an - a[n-1]
            transformed_a[n] = delta*last
        else:
            last = sin(an)
    return np.asarray(transformed_a) * b
Run Code Online (Sandbox Code Playgroud)

速度快了 10 倍多一点。

cython -a在这里很有用 - 它会生成一个带注释的 HTML 文件,显示哪些行包含与 Python 的大量交互。