Cython 优化速度慢

Question

我正在尝试使用 cython 优化以下 python 代码：

from cython cimport boundscheck, wraparound

@boundscheck(False)
@wraparound(False)
def cython_color2gray(numpy.ndarray[numpy.uint8_t, ndim=3] image):
    cdef int x,y,z
    cdef double z_val, grey
    for x in range(len(image)):
        for y in range(len(image[x])):
            grey = 0
            for z in range(len(image[x][y])):
                if z == 0:
                    z_val = image[x][y][0] * 0.21
                    grey += z_val
                elif z == 1:
                    z_val = image[x][y][1] * 0.07
                    grey += z_val
                elif z == 2:
                    z_val = image[x][y][2] * 0.72
                    grey += z_val
            image[x][y][0] = grey
            image[x][y][1] = grey
            image[x][y][2] = grey
    return image

但是，当检查一切是否都达到应有的优化时，我收到以下黄线（见图）。我还能做些什么来优化这个 cython 代码并使其运行得更快吗？

Answer 1

以下是一些要点：

• len()

  函数是一个Python函数，具有可测量的开销。由于

  image

  无论如何都是

  np.ndarray

  ，因此更喜欢使用

  .shape

  属性来获取每个维度中的元素数量。

• 考虑使用

  image[i, j, k]

  而不是

  image[i][j][k]

  进行元素访问。

• 更喜欢类型化内存视图，因为语法更清晰并且速度更快。例如，

  numpy.ndarray[T, ndim=3]

  的等效内存视图是

  T[:, :, :]

  ，其中

  T

  表示数据元素的类型。如果您知道阵列的内存布局是 C-连续，则可以使用

  T[:, :, ::1]

  指定布局。在 C 中，

  unsigned char

  是具有 8 位宽度的最小无符号整数类型（在大多数现代平台上），因此相当于

  np.uint8_t

  。因此，鉴于

  numpy.ndarray[numpy.uint8_t, ndim=3] image

  的数据是 C 连续的，您的

  unsigned char[:, :, ::1] image

  变为

  image

  。或者，您可以在从

  uint8_t[:, :, ::1]

  cimport

  ing C 类型

  uint8_t

  后使用

  libc.stdint

  。

• 变量

  grey

  是双精度型，而

  image

  的元素是

  np.uint8

  （相当于unsigned char）。因此，当在 Python 中执行

  image[i,j,k]=grey

  时，

  grey

  会被转换为无符号字符，即十进制数字被截断。在 Cython 中，您必须手动进行转换。

• 在您知道代码按预期工作后，您可以使用 Cython 编译器的指令进一步加速它，例如停用边界检查和负索引（环绕）。请注意，这些是需要导入的装饰器。

你的代码片段变成：

from cython cimport boundscheck, wraparound

@boundscheck(False)
@wraparound(False)
def cython_color2gray(unsigned char[:, :, ::1] image):
    cdef int x,y,z
    cdef double z_val, grey
    for x in range(image.shape[0]):
        for y in range(image.shape[1]):
            grey = 0
            for z in range(image.shape[2]):
                if z == 0:
                    z_val = image[x, y, 0] * 0.21
                    grey += z_val
                elif z == 1:
                    z_val = image[x, y, 1] * 0.07
                    grey += z_val
                elif z == 2:
                    z_val = image[x, y, 2] * 0.72
                    grey += z_val
            image[x, y, :] = <unsigned char> grey
    return image

仔细观察，你会发现不需要最内层的循环：

from cython cimport boundscheck, wraparound

@boundscheck(False)
@wraparound(False)
def cython_color2gray(unsigned char[:, :, ::1] image):
    cdef int x, y
    for x in range(image.shape[0]):
        for y in range(image.shape[1]):
            image[x, y, :] = <unsigned char>(image[x,y,0]*0.21 + image[x,y,1]*0.07 + image[x,y,2] * 0.72)
    return image

更进一步，您可以尝试通过启用 C 编译器的自动矢量化（SIMD 意义上的）来加速 Cython 生成的 C 代码。对于 gcc/clang，您可以使用标志

-O3

和

-march=native

。对于 MSVC，它是

/O2

和

/arch:AVX2

（假设您的计算机支持 AVX2）。如果您在 jupyter 笔记本中工作，您可以通过 Cython magic 的

-c=YOURFLAG

参数传递 C 编译器标志，即

%%cython -a -f -c=-O3 -c=-march=native
# your cython code here..