我一直非常喜欢 Python。当人们提到 Python 的时候,经常会说到下面两个优点:
- 写起来方便
- 容易调用 C/C++ 的库
然而实际上,第一点是以巨慢的执行速度为代价的,而第二点也需要库本身按照 Python 的规范使用 Python API、导出相应的符号。
在天壤实习的时候,跟 Cython 打了不少交道,觉得这个工具虽然 Bug 多多,写的时候也有些用户体验不好的地方,但已经能极大提高速度和方便调用 C/C++,还是非常不错的。这里就给大家简单介绍一下 Cython(注意区别于 CPython)。Cython 可以让我们方便地:
- 用 Python 的语法混合编写 Python 和 C/C++ 代码,提升 Python 速度
- 调用 C/C++ 代码
例子:矩阵乘法
假设我们现在正在编写一个很简单的矩阵乘法代码,其中矩阵是保存在 numpy.ndarray 中。Python 代码可以这么写:
# dot_python.py
import numpy as np
def naive_dot(a, b):
if a.shape[1] != b.shape[0]:
raise ValueError('shape not matched')
n, p, m = a.shape[0], a.shape[1], b.shape[1]
c = np.zeros((n, m), dtype=np.float32)
for i in xrange(n):
for j in xrange(m):
s = 0
for k in xrange(p):
s += a[i, k] * b[k, j]
c[i, j] = s
return c
不用猜也知道这比起 C/C++ 写的要慢的不少。我们感兴趣的是,怎么用 Cython 加速这个程序。我们先上 Cython 程序代码:
# dot_cython.pyx
import numpy as np
cimport numpy as np
cimport cython
@cython.boundscheck(False)
@cython.wraparound(False)
cdef np.ndarray[np.float32_t, ndim=2] _naive_dot(np.ndarray[np.float32_t, ndim=2] a, np.ndarray[np.float32_t, ndim=2] b):
cdef np.ndarray[np.float32_t, ndim=2] c
cdef int n, p, m
cdef np.float32_t s
if a.shape[1] != b.shape[0]:
raise ValueError('shape not matched')
n, p, m = a.shape[0], a.shape[1], b.shape[1]
c = np.zeros((n, m), dtype=np.float32)
for i in xrange(n):
for j in xrange(m):
s = 0
for k in xrange(p):
s += a[i, k] * b[k, j]
c[i, j] = s
return c
def naive_dot(a, b):
return _naive_dot(a, b)
可以看到这个程序和 Python 写的几乎差不多。我们来看看不一样部分:
- Cython 程序的扩展名是 .pyx
- cimport 是 Cython 中用来引入 .pxd 文件的命令。有关 .pxd 文件,可以简单理解成 C/C++ 中用来写声明的头文件,更具体的我会在后面写到。这里引入的两个是 Cython 预置的。
- @cython.boundscheck(False) 和 @cython.wraparound(False) 两个修饰符用来关闭 Cython 的边界检查
-
Cython 的函数使用 cdef 定义,并且他可以给所有参数以及返回值指定类型。比方说,我们可以这么编写整数 min 函数:
cdef int my_min(int x, int y): return x if x <= y else y
这里 np.ndarray[np.float32_t, ndim=2] 就是一个类型名就像 int 一样,只是它比较长而且信息量比较大而已。它的意思是,这是个类型为 np.float32_t 的2维 np.ndarray。
- 在函数体内部,我们一样可以使用 cdef typename varname 这样的语法来声明变量
- 在 Python 程序中,是看不到 cdef 的函数的,所以我们这里 def naive_dot(a, b) 来调用 cdef 过的 _naive_dot 函数。