numpy - 如何使用 cython 加速 numpy？

Question

我按照说明
http://docs.cython.org/en/latest/src/tutorial/numpy.html

但是当我尝试构建自己的块时遇到了一些问题:

(代码的目的只是计算两个三角形的并集面积)

我的 .pyx 代码:

cimport cython
import numpy as np
cimport numpy as np

DTYPE = np.float
ctypedef np.float_t DTYPE_t

cpdef DTYPE_t union(np.ndarray[DTYPE_t, ndim=1] au, np.ndarray[DTYPE_t, ndim=1] bu, DTYPE_t area_intersection):
    cdef DTYPE_t area_a
    cdef DTYPE_t area_b
    cdef DTYPE_t area_union
    cdef DTYPE_t a = au[2]
    cdef DTYPE_t b = au[0]
    cdef DTYPE_t c = au[3]
    cdef DTYPE_t d = au[1]
    cdef DTYPE_t e = bu[2]
    cdef DTYPE_t f = bu[0]
    cdef DTYPE_t g = bu[3]
    cdef DTYPE_t h = bu[1]
    area_a = (a - b) * (c - d)
    area_b = (e - f) * (g - h)
    area_union = area_a + area_b - area_intersection
    return area_union

我的 .py 代码

import numpy as np
import random


def union(au, bu,area_intersection):
    area_a = (au[2] - au[0]) * (au[3] - au[1])
    area_b = (bu[2] - bu[0]) * (bu[3] - bu[1])
    area_union = area_a + area_b - area_intersection
    return area_union

我的 setup.py 文件:

from distutils.core import setup
from Cython.Build import cythonize
import numpy

setup(ext_modules = cythonize('union.pyx'),include_dirs=[numpy.get_include()])

我用下面的代码来测试cython的速度:

from union_py import union as py_speed
from union import union as cy_speed
import numpy as np
import time

np.random.seed(1)
start = time.time()
for i in range (1000000):
    in_a = np.random.rand(4)
    in_b = np.random.rand(4)
    c = cy_speed(au = in_a,bu = in_b,area_intersection = 2.1)

end = time.time()
print (end - start)

对于python速度，我只是更改了 cy_speed至 py_speed .

结果显示cython占用 2.291128158569336而python需要 2.0604214668273926 . python 版本甚至更快。我确保 cython 代码的功能(计算联合区域)是正确的。如何改进 cython 代码以加快速度？

Answer 1

DavidW 的感觉是对的:cython 必须在运行时检查传递的数组的类型，这意味着由于函数本身的操作很少而无法恢复的开销。

numpy-array 不是这个任务的最佳选择 - 正如我们将看到的，使用 cdef-classes 可以将 python 击败 10 倍。

对于我的实验，我使用了稍微不同的设置:

>>> import numpy as np
>>> a=np.random.rand(4)
>>> b=np.random.rand(4)

>>> %timeit py_union(a,b,2.1)
1.3 µs ± 51.4 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

>>> %timeit cy_union(a,b,2.1)
1.39 µs ± 11.7 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

所以cython版本确实有点慢。正如 DavidW 所指出的，这是由于 cython 的类型检查，当我们查看生成的 C 代码时，在计算第一行函数之前，必须发生以下情况:

...
__Pyx_LocalBuf_ND __pyx_pybuffernd_au;
...
if (unlikely(__Pyx_GetBufferAndValidate(&__pyx_pybuffernd_au.rcbuffer->pybuffer, (PyObject*)__pyx_v_au, &__Pyx_TypeInfo_nn___pyx_t_3foo_DTYPE_t, PyBUF_FORMAT| PyBUF_STRIDES, 1, 0, __pyx_stack) == -1)) __PYX_ERR(0, 7, __pyx_L1_error)

__Pyx_GetBufferAndValidate的定义可以找到 here ，我们可以很容易地看到，它不是免费的。

我们用两个实验来验证一下。首先减少函数中的操作次数:

%%cython
import numpy as np
cimport numpy as np

ctypedef np.float_t DTYPE_t

cpdef DTYPE_t cy_silly1(np.ndarray[DTYPE_t, ndim=1] au, np.ndarray[DTYPE_t, ndim=1] bu, DTYPE_t area_intersection):
    area_union = au[0] + bu[1] - area_intersection
    return area_union

>>> %timeit cy_silly1(a,b,2.1)
1.4 µs ± 12.3 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

我们减少了函数中的操作次数，但对执行时间没有影响，即函数的这部分不是瓶颈。

如果我们只有一个 numpy 数组要检查会发生什么？

%%cython
import numpy as np
cimport numpy as np

ctypedef np.float_t DTYPE_t

cpdef DTYPE_t cy_silly2(np.ndarray[DTYPE_t, ndim=1] au, DTYPE_t area_intersection):
    cdef DTYPE_t area_union = au[0] + au[1] - area_intersection
    return area_union

>>> %timeit cy_silly2(a,2.1)
745 ns ± 7.46 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

这次我们几乎得到了加速 2 - __Pyx_GetBufferAndValidate真的是瓶颈。

可以做什么？ Typed memory views开销略少，因为它们使用完全不同的机器:

%%cython
...
cpdef DTYPE_t cy_union_tmv(DTYPE_t[::1] au, DTYPE_t[::1] bu, DTYPE_t area_intersection):
...#the same as above

%timeit cy_union_tmv(a,b,2.1)
1.09 µs ± 3.24 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

一个更好的主意是编写一个专用的 cdef 类，它的开销要少得多:

import numpy as np
cimport numpy as np

DTYPE = np.float
ctypedef np.float_t DTYPE_t

cdef class Triangle:
   cdef DTYPE_t a
   cdef DTYPE_t b
   cdef DTYPE_t c
   cdef DTYPE_t d
   def __init__(self, a,b,c,d):
      self.a=a
      self.b=b
      self.c=c
      self.d=d
   cdef DTYPE_t get_area(self):
      return (self.a-self.b)*(self.c-self.d)


cpdef DTYPE_t cy_union_cdef(Triangle au, Triangle bu, DTYPE_t area_intersection):
    cdef DTYPE_t area_union = au.get_area() + bu.get_area() - area_intersection 
    return area_union

现在:

>>> tri_a=Triangle(a[0],a[1],a[2],a[3])
>>> tri_b=Triangle(b[0],b[1],b[2],b[3]) 
>>> %timeit cy_union_cdef(tri_a,tri_b,2.1)
106 ns ± 0.668 ns per loop (mean ± std. dev. of 7 runs, 10000000 loops each)

产生约 10 的加速。