教你学会cython加速

By 青衣极客 Blue Geek In 2020-01-04

阅读一些开源代码时，常常碰到cython这个第三方模块，特别是在图像处理、计算机视觉以及深度学习的项目中。究其原因，还是python的老问题，运行速度太慢。这么说cython就是加速python的了？确实如此，看这个模块的命名大概就能猜到其功能是C语言与python结合相关的。关于利用cython加速你的python程序将在下文揭晓。

1. 使用cython的动机

为什么要使用cython？这是个自然而然的问题。我们知道，python的语法集合已经基本足够实现各种功能，但对于运行速度慢这个问题实在是无能为力。那么究竟慢到什么程度呢？在演示例子之前先导入所需的模块，并设置对应的环境变量。其中“code/cython”是cython代码存放的路径。使用numpy生成一个大数组用于测试。

import sys
sys.path.append('code/cython/')
import numpy as np
from math import sqrt

X = np.random.rand(1920, 1920)

准备条件已经就绪，我们先来说一说演示的思路。熟悉python的朋友大概知道，python运行效率低的原因大概在两点：1. 数据结构臃肿；2. 循环运行慢。这里选择对一个矩阵中的每个元素进行平方根的计算，既包括数据结构，也包括较大的循环，能够说明一般python 程序的特点。示例函数编写完成后，通过notebook的magic工具timeit实现对函数运行耗时的测试。

def psqrt(X):
    Y = np.zeros(X.shape)
    for i in range(X.shape[0]):
        for j in range(X.shape[1]):
            Y[i, j] = sqrt(X[i, j])
    return Y

%timeit psqrt(X)

1.23 s ± 47.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

从测试结果来看，处理一次1920x1920的矩阵居然耗时达到1秒多。通常我们可能需要处理上万个这种数据，那么所耗时间将很难容忍。遇到这种情况，那就需要寻找加速的办法了。首先，我们当然可以使用C扩展的方式改写这部分功能，但是C扩展操作起来着实比较麻烦。如果又想加速，又想简单，那么有没有办法呢？这就该cython登场了。因此，使用cython的动机是你对程序运行速度慢这件事已经没有容忍度了，但还保有程序员的优良品质——懒惰。

2. 使用cython的方法

使用cython这件事已经确凿无疑，接下来就是讲述操作流程了。第一步，需要在pyx文件中使用cython的语法规范实现功能。一种新的语法规范听起来比较吓人，但实际上，这种规范就是包含C语言数据结构的python代码。代码编写的语句语法为python，只是其中特定的数据结构声明为C的数据结构即可。如果还是不明白，可以看看下面的cython代码示例，实现的功能与上文的python代码一致。

!cat code/cython/csqrt.pyx

# cython: language_level=3
from math import sqrt
import numpy as np

def csqrt(double[:,:] X):
    cdef int r = X.shape[0]
    cdef int c = X.shape[1]
    cdef double[:, :] Y = np.zeros((r, c))
    cdef int i, j

    for i in range(r):
        for j in range(c):
            Y[i, j] = sqrt(X[i, j])
    return Y  

cython代码编写完成之后，就可以开始第二步：调用pyx代码。第一种调用方式是编写setup.py文件或者手动使用编译器将pyx文件编译成动态链接库so或者dll文件。这种方式与一般的C扩展是一样的，如果选择这样做那倒不如直接使用C扩展。第二种调用方式是使用cython模块自带的pyximport工具直接在python代码中加载pyx模块。这里就是使用第二种方式来调用pyx代码。

import pyximport;pyximport.install()
import csqrt

%timeit csqrt.csqrt(X)

107 ms ± 214 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

经过耗时测试之后发现，相同的功能，cython代码只需要100多毫秒就可以完成，相当于速度提升了10倍，这在程序速度优化上可是了不得的成绩。但是100毫秒耗时仍然有些不太能接受，那么还有没有办法继续优化呢？

3. 性能改进

要想知道有没有继续优化的空间，那么还是回到最初的问题，明确python程序运行效率低的原因。上面的cython代码中数据结构已经换成了C语言的数据结构，剩下就该看看循环中的内容了。经过反复检查，我们发现，上面cython代码在循环中调用了一个python函数sqrt，如果将这个函数换成C语言库中的调用，那么速度应该会有进一步的提升。相对于上面的cython代码，只需要将sqrt的来源更换一下即可。

!cat code/cython/csqrt2.pyx

# cython: language_level=3
from libc.math cimport sqrt
import numpy as np

def csqrt(double[:,:] X):
    cdef int r = X.shape[0]
    cdef int c = X.shape[1]
    cdef double[:, :] Y = np.zeros((r, c))
    cdef int i, j

    for i in range(r):
        for j in range(c):
            Y[i, j] = sqrt(X[i, j])
    return Y  

其他代码都不变，只是将sqrt的来源从python的math模块换成了C语言的math库。接下来的测试将见证奇迹。

import csqrt2
%timeit csqrt2.csqrt(X)

10.6 ms ± 65.8 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

从测试结果来看，相同的功能在优化性能之后只需大约10毫秒就可以完成。相对于最开始的python代码，最新的cython程序运行速度提升了大约100倍。这种速度的提升还真是令人难以置信，不过事实就摆在眼前。我们是不是该为cython对一般python程序速度的优化而欢呼，是不是该把之前编写的python代码全都使用cython优化一遍，是不是该在开发新项目时直接使用cython？

4. 使用cython的时机

cython虽然神奇，但也不要太过夸大。在上面演示的这样功能上，如果我们调用numpy封装的sqrt函数，那么耗时又是怎样的呢？

%timeit np.sqrt(X)

6.31 ms ± 23.1 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

测试结果为大约6毫秒，比我们上面花了老大劲优化的cython代码还快。看到这种结果大概是欲哭无泪，要跟cython说拜拜了。不过cython可并不是吃干饭，那么究竟在哪些情况下该使用cython呢？一般而言，我们使用python是因为其简单易用，开发效率高。而在一段程序中拖累运行速度的通常只是一小部分功能代码。如果能够将那些耗时90%的小部分代码改用cython实现，那么性能优化和开发效率就可以兼顾了。因此，通常开发仍然是使用python完成的，只有发现运行速度慢到不可接受时才考虑使用cython。使用cython也并不是全面铺开，而是选择耗时最大的几个功能点进行优化。此外，还有一点需要牢记，第三方库已经优化过的基本操作不需要用cython重新实现，因为优化收益已经不大了，本文计算矩阵元素平方根的例子就是明证。这就是使用cython的时机。掌握了cython的使用方法和时机，再也不怕python运行慢了。

• Previous
  graphviz让你理清复杂的关系
• Next
  深度学习帮你识别验证码